Company Queries II

def dfs(v, p):
    up[0][v] = p
    for u in adj[v]:
        if u != p:
            depth[u] = depth[v] + 1
            dfs(u, v)


def lca(a, b):
    if depth[a] < depth[b]:
        a, b = b, a
    # Lift a up to depth[b]
    diff = depth[a] - depth[b]
    for i in range(LOG):
        if diff & (1 << i):
            a = up[i][a]

    if a == b:
        return a

    # Lift a and b up until their parents are equal
    for i in reversed(range(LOG)):
        if up[i][a] != up[i][b]:
            a = up[i][a]
            b = up[i][b]

    return up[0][a]


if __name__ == "__main__":

    n, q = map(int, input().split())
    bosses = list(map(int, input().split()))

    queries = []

    for _ in range(q):
        a, b = map(int, input().split())
        queries.append([a, b])

    # Build adjacency list for the tree
    adj = [[] for _ in range(n + 1)]
    for i, b in enumerate(bosses, start=2):
        adj[b].append(i)

    LOG = 20

    up = [[0] * (n + 1) for _ in range(LOG)]
    depth = [0] * (n + 1)

    dfs(1, 0)

    print(adj)
    print(depth)

    # Precompute all 2^j ancestors
    for j in range(1, LOG):
        for v in range(1, n + 1):
            up[j][v] = up[j - 1][up[j - 1][v]]

    for i in range(q):
        print(lca(queries[i][0], queries[i][1]))

Test 1

Verdict: WRONG ANSWER

Test 2

Verdict: WRONG ANSWER

Test 3

Verdict: WRONG ANSWER

Test 4

Verdict: WRONG ANSWER

Test 5

Verdict: WRONG ANSWER

Test 6

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 51, in <module>
    dfs(1, 0)
  File "input/code.py", line 6, in dfs
    dfs(u, v)
  File "input/code.py", line 6, in dfs
    dfs(u, v)
  File "input/code.py", line 6, in dfs
    dfs(u, v)
  [Previous line repeated 1576 more times]
  File "input/code.py", line 2, in dfs
    up[0][v] = p
RecursionError: maximum recursion depth exceeded

Test 7

Verdict: WRONG ANSWER

Test 8

Verdict: TIME LIMIT EXCEEDED

Test 9

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 51, in <module>
    dfs(1, 0)
  File "input/code.py", line 6, in dfs
    dfs(u, v)
  File "input/code.py", line 6, in dfs
    dfs(u, v)
  File "input/code.py", line 6, in dfs
    dfs(u, v)
  [Previous line repeated 1582 more times]
  File "input/code.py", line 2, in dfs
    up[0][v] = p
RecursionError: maximum recursion depth exceeded

Test 10

Verdict: TIME LIMIT EXCEEDED

Test 11

Verdict: WRONG ANSWER

Test 12

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 51, in <module>
    dfs(1, 0)
  File "input/code.py", line 6, in dfs
    dfs(u, v)
  File "input/code.py", line 6, in dfs
    dfs(u, v)
  File "input/code.py", line 6, in dfs
    dfs(u, v)
  [Previous line repeated 1576 more times]
  File "input/code.py", line 2, in dfs
    up[0][v] = p
RecursionError: maximum recursion depth exceeded