Cod sursa(job #2540119)

Utilizator TeodorCotetCotet Teodor TeodorCotet Data 6 februarie 2020 19:18:38
Problema Lowest Common Ancestor Scor 30
Compilator py Status done
Runda Arhiva educationala Marime 1.43 kb
MAX_DIM = 18

def read_gen(fname):
    with open(fname, 'rt') as fin:
        for line in fin:
            for val in line.split():
                yield int(val)


def build_ancestors(anc, n):
    global MAX_DIM
    for p in range(1, MAX_DIM):
        for i in range(1, n + 1):
            anc[i][p] = anc[anc[i][p - 1]][p - 1]
        
def dfs(node, sons, lvl, current_lvl):
    lvl[node] = current_lvl
    for x in sons[node]:
        dfs(x, sons, lvl, current_lvl + 1)

def lca(x, y, anc, lvl):
    if lvl[x] < lvl[y]:
        x, y = y, x

    for i in range(MAX_DIM - 1, -1, -1):
        if lvl[anc[x][i]] >= lvl[y]:
            x = anc[x][i]
    if x == y:
        return x

    for i in range(MAX_DIM - 1, -1, -1):
        if anc[x][i] > 0  and anc[y][i] > 0 and anc[x][i] != anc[y][i]:
            x = anc[x][i]
            y = anc[y][i]

    return anc[x][0] 

if __name__ == "__main__":

    it = read_gen('lca.in')

    n, m = next(it), next(it)
    sons = {x: [] for x in range(1, n + 1)}
    anc =  {i: {j: 0 for j in range(0, MAX_DIM)} for i in range(0, n + 1)}
    lvl = {i: 0 for i in range(0, n + 1)}

    for i in range(2, n + 1):
        x = next(it)
        sons[x].append(i)
        anc[i][0] = x
    dfs(1, sons, lvl, 1)
    build_ancestors(anc, n) 
    
    with open('lca.out', 'wt') as fout:
       for _ in range(m):
            x, y = next(it), next(it)
            cmn_anc = lca(x, y, anc, lvl)
            fout.write('{}\n'.format(cmn_anc))