[Python]최소공통조상 백준 13511

 문제

 https://www.acmicpc.net/problem/13511

앞의 문제들과 동일하게 DP를 이용해 2^k 번째 부모노드와 그 부모노드까지의 거리를 

구하고 최소공통조상 노드를 구하여 각 쿼리들을 해결하면된다.

u->v로 가는 비용은  u->lca까지의 거리와 v->lca까지의 거리를 구해서 더하여 구하였고

u->v로 가는 경로에 존재하는 정점중 k번째 노드는 u->lca까지의 노드개수가 k보다 

크다면 u의k번째 부모노드를 구하면되고

그렇지 않다면 k에서 u->lca  까지의 노드 개수를 빼주고 이값을 v->lca에서 뺀 값 k2라고 

했을때 v의 k2번째 부모노드를 구하면된다.

논리적으로는 틀린것이 없다고 생각하였는데 구현해보니 틀렸다.

예외적인 부분들을 몇몇 수정해 봤지만 계속 틀렸다. 테스트 케이스도 바꾸어가며 실험해 

봤지만 틀린걸 찾을수 없었다.혹시 이 소스의 문제점이나 반례를 가르쳐 주실수 있다면 

댓글을 통해 알려주시면 감사하겠습니다.

여러가지 예외를 처리하면서 막 구현하다 보니 소스가 지저분 할 수 있습니다.

import sys
from math import log2
from collections import deque
 
N=int(sys.stdin.readline())
logN=int(log2(N)+1)
tree=[[] for _ in range(N+1)]
for _ in range(N-1):
    u,v,w=map(int,sys.stdin.readline().split())
    tree[u].append([v,w])
    tree[v].append([u,w])
 
 
p_list=[[0,0] for _ in range(N+1)]
depth=[0 for _ in range(N+1)]
p_check=[True for _ in range(N+1)]
#부모노드 저장
q=deque()
q.append(1)
while q:
    a=q.popleft()
    p_check[a]=False
    for b,c in tree[a]:
        if p_check[b]:
            p_list[b][0]=a#부모노드 저장
            p_list[b][1]=c#부모노드까지의 거리 저장
            q.append(b)
            depth[b]=depth[a]+1
 
 
#2^k번째 부모 노드와 2^k번째 부모 노드까지의 거리
DP=[[[0,0] for _ in range(logN)] for _ in range(N+1)]#2^k번째의 부모노드, 2^k번째 부모 노드까지의 거리
#초기화
for i in range(N+1):
    DP[i][0][0]=p_list[i][0]
    DP[i][0][1]=p_list[i][1]
 
 
#희소테이블 완성하기
for j in range(1,logN):
    for i in range(1,N+1):
        DP[i][j][0]=DP[DP[i][j-1][0]][j-1][0]
        if DP[i][j][0]!=0:
            DP[i][j][1] = DP[i][j - 1][1] + DP[DP[i][j - 1][0]][j - 1][1]
 
 
M=int(sys.stdin.readline())
for _ in range(M):
    Q=list(map(int,sys.stdin.readline().split()))
    # 깊이차이
    a=Q[1]
    b=Q[2]
    if depth[a] < depth[b]:
        a, b = b, a
 
    dif = depth[a] - depth[b]
    # 레벨맞추기
    for i in range(logN):
        if dif & 1 << i:
            a = DP[a][i][0]
    if a==b:
        LCA=a
    else:
    #최소공통조상 찾기
        for i in range(logN-1,-1,-1):
            if DP[a][i][0]!=DP[b][i][0]:
                a=DP[a][i][0]
                b=DP[b][i][0]
        #최소공통조상
        LCA = DP[a][0][0]
    #최소공통조상의 레벨
    lca_depth=depth[LCA]
 
    if Q[0]==1:
        sum = 0
        #각 노드에서 최소공통조상까지 거리 구해서 더하기
        dif_a=depth[Q[1]]-lca_depth
        dif_b=depth[Q[2]]-lca_depth
        # print(dif_a)
        # print(dif_b)
        for i in range(logN):
            if dif_a & 1<<i:
                sum +=DP[Q[1]][i][1]
                Q[1]=DP[Q[1]][i][0]
 
            if dif_b & 1<<i:
                sum +=DP[Q[2]][i][1]
                Q[2]=DP[Q[2]][i][0]
        print(sum)
 
 
 
    elif Q[0]==2:
        #Q[3]이 Q[1]과 CA 사이의 노드수보다 높을거나 같을때는 Q[3]에서 Q[1]~LCA를 뺀 나머지를 통해 구함
        #낮을때는 그냥 구함
        #한 직선상에 잇고 Q[1]이 레벨이 더 낮을때
        if LCA==Q[1]:
            if Q[3]==1:
                print(Q[1])
            else:
                gep=depth[Q[2]]-Q[3]
                for i in range(logN):
                    if gep & 1<<i:
                        Q[2]=DP[Q[2]][i][0]
                print(Q[2])
            continue
 
 
        else:
            gep=depth[Q[1]]-lca_depth+1
            if Q[3]<=gep:
                for i in range(logN):
                    if Q[3]-1 & 1<<i:
                        Q[1]=DP[Q[1]][i][0]
                print(Q[1])
 
            else:
                # Q[3]=(depth[Q[2]]-lca_depth)-(Q[3]-(depth[Q[1]-lca_depth]))
                Q[3]=depth[Q[2]]-Q[3]+depth[Q[1]]-2*lca_depth+1
                for i in range(logN):
                    if Q[3] & 1<<i:
                        Q[2]=DP[Q[2]][i][0]
                print(Q[2])

문제는 예외 처리부분에 있었습니다.

#한 직선상에 잇고 Q[1]이 레벨이 더 낮을때
 if LCA==Q[1]:
     if Q[3]==1:
         print(Q[1])
     else:
         gep=depth[Q[2]]-Q[3]
         for i in range(logN):
             if gep & 1<<i:
                 Q[2]=DP[Q[2]][i][0]
         print(Q[2])
     continue
 else:

예외적인 부분이라고 생각했으나 그냥 아래의 코드로 처리 하면된다.

import sys
from math import log2
from collections import deque
 
N=int(sys.stdin.readline())
logN=int(log2(N)+1)
tree=[[] for _ in range(N+1)]
for _ in range(N-1):
    u,v,w=map(int,sys.stdin.readline().split())
    tree[u].append([v,w])
    tree[v].append([u,w])
 
 
p_list=[[0,0] for _ in range(N+1)]
depth=[0 for _ in range(N+1)]
p_check=[True for _ in range(N+1)]
#부모노드 저장
q=deque()
q.append(1)
while q:
    a=q.popleft()
    p_check[a]=False
    for b,c in tree[a]:
        if p_check[b]:
            p_list[b][0]=a#부모노드 저장
            p_list[b][1]=c#부모노드까지의 거리 저장
            q.append(b)
            depth[b]=depth[a]+1
 
 
#2^k번째 부모 노드와 2^k번째 부모 노드까지의 거리
DP=[[[0,0] for _ in range(logN)] for _ in range(N+1)]
#초기화
for i in range(N+1):
    DP[i][0][0]=p_list[i][0]
    DP[i][0][1]=p_list[i][1]
 
 
#희소테이블 완성하기
for j in range(1,logN):
    for i in range(1,N+1):
        DP[i][j][0]=DP[DP[i][j-1][0]][j-1][0]
        if DP[i][j][0]!=0:
            DP[i][j][1] = DP[i][j - 1][1] + DP[DP[i][j - 1][0]][j - 1][1]
 
 
M=int(sys.stdin.readline())
for _ in range(M):
    Q=list(map(int,sys.stdin.readline().split()))
    # 깊이차이
    a=Q[1]
    b=Q[2]
    if depth[a] < depth[b]:
        a, b = b, a
 
    dif = depth[a] - depth[b]
    # 레벨맞추기
    for i in range(logN):
        if dif & 1 << i:
            a = DP[a][i][0]
    if a==b:
        LCA=a
    else:
    #최소공통조상 찾기
        for i in range(logN-1,-1,-1):
            if DP[a][i][0]!=DP[b][i][0]:
                a=DP[a][i][0]
                b=DP[b][i][0]
        #최소공통조상
        LCA = DP[a][0][0]
    #최소공통조상의 레벨
    lca_depth=depth[LCA]
 
    if Q[0]==1:
        sum = 0
        #각 노드에서 최소공통조상까지 거리 구해서 더하기
        dif_a=depth[Q[1]]-lca_depth
        dif_b=depth[Q[2]]-lca_depth
        # print(dif_a)
        # print(dif_b)
        for i in range(logN):
            if dif_a & 1<<i:
                sum +=DP[Q[1]][i][1]
                Q[1]=DP[Q[1]][i][0]
 
            if dif_b & 1<<i:
                sum +=DP[Q[2]][i][1]
                Q[2]=DP[Q[2]][i][0]
        print(sum)
 
 
 
    elif Q[0]==2:
        #Q[3]이 Q[1]과 CA 사이의 노드수보다 높을거나 같을때는 Q[3]에서 Q[1]~LCA를 뺀 나머지를 통해 구함
        #낮을때는 그냥 구함
        #한 직선상에 잇고 Q[1]이 레벨이 더 낮을때
        # if LCA==Q[1]:
        #     if Q[3]==1:
        #         print(Q[1])
        #     else:
        #         gep=depth[Q[2]]-Q[3]
        #         for i in range(logN):
        #             if gep & 1<<i:
        #                 Q[2]=DP[Q[2]][i][0]
        #         print(Q[2])
        #     continue
        # else:
        gep=depth[Q[1]]-lca_depth+1
        if Q[3]<=gep:
            for i in range(logN):
                if Q[3]-1 & 1<<i:
                    Q[1]=DP[Q[1]][i][0]
            print(Q[1])
 
        else:
            # Q[3]=(depth[Q[2]]-lca_depth)-(Q[3]-(depth[Q[1]-lca_depth]))
            Q[3]=depth[Q[2]]-Q[3]+depth[Q[1]]-2*lca_depth+1
            for i in range(logN):
                if Q[3] & 1<<i:
                    Q[2]=DP[Q[2]][i][0]
            print(Q[2])
 
 

이렇게 예외처리부분을 주석처리 해주면 된다. pypy3로 통과가 가능하였고 python3으

로는 시간초과가 났다.속도를 조금 높이는 방법이 있다면 DP에 2^k번째 부모노드까지의

 거리를 구하는대신 각 노드에서 루트 노드까지의 거리를 저장하여

 u->root+v->root-(lca*2)를 통해 u->v까지의 비용을 구하는 방법이 있다. 

공간복잡도 측면에서도 더욱 효율적이다.