fork download
copy 
  1. from collections import defaultdict, deque
  2.  
  3. def find_max_length(t, test_cases):
  4.     results = []
  5.  
  6.     for n, a in test_cases:
  7.         b = [a[i] + i for i in range(n)]  # b[i] = a[i] + (i + 1)
  8.         c = [b[i] + i for i in range(n)]  # c[i] = b[i] + (i + 1)
  9.  
  10.         u = [(b[i], i) for i in range(n)]  # Pair (b[i], i)
  11.         v = [(c[i], i) for i in range(n)]  # Pair (c[i], i)
  12.  
  13.         # Sort both u and v based on the first element of the pair
  14.         u.sort()
  15.         v.sort()
  16.  
  17.         # Use two pointers to find pairs (i, j) such that u[i].first == v[j].first
  18.         graph = defaultdict(list)
  19.  
  20.         i, j = 0, 0
  21.         while i < n and j < n:
  22.             if u[i][0] == v[j][0]:  # Found a match
  23.                 graph[u[i][1]].append(v[j][1])
  24.                 i += 1
  25.             elif u[i][0] < v[j][0]:
  26.                 i += 1
  27.             else:
  28.                 j += 1
  29.  
  30.         # Initialize d array and perform BFS/DFS to find reachable nodes
  31.         d = [0] * n
  32.         visited = [False] * n
  33.  
  34.         def bfs(start):
  35.             queue = deque([start])
  36.             reachable_values = []
  37.             while queue:
  38.                 node = queue.popleft()
  39.                 if visited[node]:
  40.                     continue
  41.                 visited[node] = True
  42.                 reachable_values.append(c[node])
  43.                 for neighbor in graph[node]:
  44.                     if not visited[neighbor]:
  45.                         queue.append(neighbor)
  46.             return max(reachable_values, default=0)
  47.  
  48.         # Compute d[i] for each vertex
  49.         for i in range(n):
  50.             if not visited[i]:
  51.                 max_c = bfs(i)
  52.                 d[i] = max_c
  53.  
  54.         count = 0
  55.         # Count valid positions
  56.         for i in range(n):
  57.             if b[i] == n:
  58.                 count = max(count, d[i])
  59.  
  60.         # The answer for this test case
  61.         results.append(count)
  62.  
  63.     return results
  64.  
  65. # Read input
  66. t = int(input())
  67. test_cases = []
  68.  
  69. for _ in range(t):
  70.     n = int(input())
  71.     a = list(map(int, input().split()))
  72.     test_cases.append((n, a))
  73.  
  74. # Get results
  75. results = find_max_length(t, test_cases)
  76.  
  77. # Print results
  78. for res in results:
  79.     print(res)
  80.  
Success #stdin #stdout 0.04s 9960KB
comments (?)
stdin
copy 
4
5
2 4 6 2 5
5
5 4 4 5 1
4
6 8 2 3
1
1
stdout
copy 
6
9
6
1
https://ideone.com/8TFGoC
language:
Python 3 (python  3.9.5)
created:
5 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!