Company Queries II

import sys
import sys
def main():
    import sys
    import sys
    from math import log2, ceil
    from sys import stdin
 
    data = sys.stdin.read().split()
    ptr = 0
 
    n = int(data[ptr]); ptr += 1
    q = int(data[ptr]); ptr += 1
 
    # Initialize adjacency list
    adj = [[] for _ in range(n + 1)]
 
    # Read bosses for employees 2 to n
    for child in range(2, n + 1):
        boss = int(data[ptr]); ptr +=1
        adj[boss].append(child)
 
    LOG = ceil(log2(n)) if n >1 else 1
    up = [[-1]*(n +1) for _ in range(LOG)]
    depth = [0]*(n +1)
 
    # BFS to compute depth and up[0][v]
    queue = [1]
    head = 0
    up[0][1] = -1  # root has no parent
 
    while head < len(queue):
        node = queue[head]
        head +=1
        for child in adj[node]:
            depth[child] = depth[node] +1
            up[0][child] = node
            queue.append(child)
 
    # Precompute up tables
    for k in range(1, LOG):
        for v in range(1, n +1):
            if up[k-1][v] != -1:
                up[k][v] = up[k-1][up[k-1][v]]
            else:
                up[k][v] = -1
 
    # Function to find LCA
    def lca(a, b):
        if depth[a] < depth[b]:
            a, b = b, a
        # Bring a to the same depth as b
        for k in range(LOG-1, -1, -1):
            if up[k][a] != -1 and depth[up[k][a]] >= depth[b]:
                a = up[k][a]
        if a == b:
            return a
        # Now lift both a and b
        for k in range(LOG-1, -1, -1):
            if up[k][a] != -1 and up[k][a] != up[k][b]:
                a = up[k][a]
                b = up[k][b]
        return up[0][a]
 
    output = []
    for _ in range(q):
        a = int(data[ptr]); ptr +=1
        b = int(data[ptr]); ptr +=1
        output.append(str(lca(a,b)))
 
    print('\n'.join(output))
 
if __name__ == "__main__":
    main()

Test 1

Verdict: ACCEPTED

Test 2

Verdict: ACCEPTED

Test 3

Verdict: ACCEPTED

Test 4

Verdict: ACCEPTED

Test 5

Verdict: ACCEPTED

Test 6

Verdict: TIME LIMIT EXCEEDED

Test 7

Verdict: TIME LIMIT EXCEEDED

Test 8

Verdict: TIME LIMIT EXCEEDED

Test 9

Verdict: TIME LIMIT EXCEEDED

Test 10

Verdict: TIME LIMIT EXCEEDED

Test 11

Verdict: ACCEPTED

Test 12

Verdict: TIME LIMIT EXCEEDED