Company Queries II

import sys

## Find lowest common ancestor in tree.
#https://www.geeksforgeeks.org/dsa/cses-solutions-company-queries-ii/

sys.setrecursionlimit(300000)

# Lowest common ancestor
LCA = []
# Depth of nodes
DEPTH = []
# Parent to child
TREE = []

LOG = 18

## Find depth of all nodes and initialize LCA
def dfs(a, b):
    # First LCA is b
    LCA[a][0] = b

    # Fill LCA. Ancestor DP table.
    for i in range(1, LOG):
        if LCA[a][i - 1] != -1:
            LCA[a][i] = LCA[LCA[a][i - 1]][i - 1]
        else:
            LCA[a][i] = -1

    # Explore children. Set same depth + 1.
    for u in TREE[a]:
        if u != b:
            DEPTH[u] = DEPTH[a] + 1
            dfs(u, a)

# Find LCA of a and b
def lca(a, b):
    if DEPTH[a] < DEPTH[b]:
        a, b = b, a

    # What is the difference in depth between a and b.
    diff = DEPTH[a] - DEPTH[b]
    for i in range(LOG):
        if (diff >> i) & 1:
            a = LCA[a][i]

    # Check if we have reached the parent
    if a == b:
        return a

    # Same level, but in a differnet branch
    for i in range(LOG - 1, -1, -1):
        if LCA[a][i] != LCA[b][i]:
            a = LCA[a][i]
            b = LCA[b][i]

    return LCA[a][0]


def main():
    global TREE, LCA, DEPTH
    
    read = input().split()
    n = int(read[0])
    q = int(read[1])

    e = (int(x) for x in input().split())
    queries = [[0, 0]] * q
    for query in range(q):
        queries[query] = [int(x) for x in input().split()]

    # Build tree
    TREE = [[] for _ in range(n + 1)]
    for child, boss in enumerate(e, start=2):
        TREE[boss].append(child)
        TREE[child].append(boss)

    # Lowest common ancestor table size
    LCA = [[-1] * LOG for _ in range(n + 1)]
    # Depth for each node
    DEPTH = [0] * (n + 1)

    dfs(1, -1)

    output = []
    for query in queries:
        a = query[0]
        b = query[1]
        output.append(str(lca(a, b)))

    for out in output:
        print(out)

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: ACCEPTED

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