PWALK - Dạo chơi đồng cỏ

Tác giả: khuc_tuan

Ngôn ngữ: Python

import gc  
gc.disable()  
import sys
import psyco
psyco.full()

[n,q] = [int(x) for x  in raw_input().split()]
ke = [[] for i in range(n+1)]
question = [[] for i in range(q)]
lq = [[] for i in range(n+1)]
vs = [False for i in range(n+1)]
H = [0 for i in range(n+1)]
fa= [0 for i in range(n+1)]

F = [-1 for i in range(n+1)]
L = [i for i in range(n+1)]

def getroot(i):
    if F[i]<0: return i
    else: return getroot(F[i])
    
def merge(i,j):
    i = getroot(i)
    j = getroot(j)
    label = L[i]
    
    if F[i]<F[j]:
        F[i] += F[j]
        F[j] = i
        L[i] = label
    else:
        F[j] += F[i]
        F[i] = j
        L[j] = label

def dfs(i):
    vs[i] = True
    
    # answer question here
    for q in lq[i]:
        u = question[q][0]
        v = question[q][1]
        if vs[u] and vs[v]:
            k = u + v - i
            LCA = L[getroot(k)]
            question[q].append(LCA)
    
    for kk in ke[i]:
        if not vs[kk[0]]:
            H[kk[0]] = H[i] + kk[1]
            fa[kk[0]] = i
            dfs(kk[0])
            
    # merge here
    if fa[i]>0:
        merge(fa[i], i)

def main():
    for i in range(n-1):
        [a,b,c] = [int(x) for x in raw_input().split()]
        ke[a].append([b,c])
        ke[b].append([a,c])
    for i in range(q):
        question[i] = [int(x) for x in raw_input().split()]
        lq[question[i][0]].append(i)
        lq[question[i][1]].append(i)
    dfs(1)
    
    for z in question:
        u = z[0]
        v = z[1]
        lca = z[2]
        
        result = H[u] + H[v] - 2 * H[lca]
        
        print result

main()

Download