Company Queries II

import heapq

n, q  = [int(x) for x in input().split()]
boss = [int(x) for x in input().split()]
queries = []
for _ in range(q):
    a, b = [int(x) for x in input().split()]
    queries.append([a-1, b-1])

treeFromDown = [0 for _ in range(n)]
treeFromUp = [[] for _ in range(n)]
for i in range(1, n):
    treeFromDown[i] = boss[i-1]-1
    treeFromUp[boss[i-1]-1].append(i)

def howManyLevelUp(x, y):
    levelX = 0
    xFound = False
    levelY = 0
    yFound = False
    heap = [(0,0)]

    while not xFound or not yFound:
        node, level = heapq.heappop(heap)
        if node == x:
            levelX = level
            xFound = True
        elif node == y:
            levelY = level
            yFound = True
        
        for child in treeFromUp[node]:
            heapq.heappush(heap, (child, level+1))
    
    return levelX, levelY

for query in queries:
    a, b = query[0], query[1]
    levelA, levelB = howManyLevelUp(a, b)
    if levelA > levelB:
        for _ in range(levelA-levelB):
            b = treeFromDown[b]
    elif levelB > levelA:
        for _ in range(levelB-levelA):
            a = treeFromDown[a]

    while a != b:
        a = treeFromDown[a]
        b = treeFromDown[b]
    
    print(a+1)

Test 1

Verdict: WRONG ANSWER

Test 2

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 39, in <module>
    levelA, levelB = howManyLevelUp(a, b)
  File "input/code.py", line 24, in howManyLevelUp
    node, level = heapq.heappop(heap)
  File "/usr/lib/pypy3.8/heapq.py", line 137, in heappop
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
IndexError: pop from empty list

Test 3

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 39, in <module>
    levelA, levelB = howManyLevelUp(a, b)
  File "input/code.py", line 24, in howManyLevelUp
    node, level = heapq.heappop(heap)
  File "/usr/lib/pypy3.8/heapq.py", line 137, in heappop
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
IndexError: pop from empty list

Test 4

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 39, in <module>
    levelA, levelB = howManyLevelUp(a, b)
  File "input/code.py", line 24, in howManyLevelUp
    node, level = heapq.heappop(heap)
  File "/usr/lib/pypy3.8/heapq.py", line 137, in heappop
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
IndexError: pop from empty list

Test 5

Verdict: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 39, in <module>
    levelA, levelB = howManyLevelUp(a, b)
  File "input/code.py", line 24, in howManyLevelUp
    node, level = heapq.heappop(heap)
  File "/usr/lib/pypy3.8/heapq.py", line 137, in heappop
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
IndexError: pop from empty list

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: RUNTIME ERROR

Traceback (most recent call last):
  File "input/code.py", line 39, in <module>
    levelA, levelB = howManyLevelUp(a, b)
  File "input/code.py", line 24, in howManyLevelUp
    node, level = heapq.heappop(heap)
  File "/usr/lib/pypy3.8/heapq.py", line 137, in heappop
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
IndexError: pop from empty list

Test 12

Verdict: TIME LIMIT EXCEEDED