[BOJ] 11812 - K진 트리 (Python)

 

 

아이디어

 

최소 공통 조상을 찾을 때까지 번호가 큰 노드를 부모 노드로 이동한다.

 

 

풀이

 

K가 1인 경우 두 노드의 거리는 노드 번호의 차이이므로 x와 y의 차이를 출력한다.

K가 2 이상인 경우 함수 find을 호출하여 두 노드의 거리를 출력한다.

함수 find는 while문에서 두 노드 중 번호가 큰 노드를 부모 노드로 이동하고, 최소 공통 조상을 찾는다.

최소 공통 조상을 찾는 동안 while문이 돌면서 거리를 증가시키고, 그 거리를 반환한다.

 

(노드 번호-2)//K-1의 식을 통해 해당 노드의 부모 노드 번호를 계산할 수 있다.

 

import sys

input = sys.stdin.readline

def find(x, y):
    d = 0
    while x != y:
        if x < y:
            y = (y-2)//K+1
        else:
            x = (x-2)//K+1
        d += 1
    return d

N, K, Q = map(int, input().split())
if K == 1:
    for _ in range(Q):
        x, y = map(int, input().split())
        print(abs(y-x))
else:
    for _ in range(Q):
        x, y = map(int, input().split())
        print(find(x, y))

 

 

for문에서 while문을 도는 것보다 함수로 빼서 돌리는 코드가 더 빨랐다.

# 더 느린 코드
for _ in range(Q):
    x, y = map(int, input().split())
    res = 0
    while x != y:
        if x < y:
            y = (y-2)//K+1
        else:
            x = (x-2)//K+1
        res += 1
    print(res)

 

처음에 두 노드의 깊이를 맞춰줘야 한다고 생각하여, K의 거듭제곱을 누적합으로 각 노드의 깊이를 찾도록 했었다.

그러나 부모 노드의 번호가 무조건 더 작으므로 노드의 깊이를 구할 필요가 없었다.

# 더 느린 코드
for _ in range(Q):
    x, y = map(int, input().split())
    if x > y:
        x, y = y, x
    acc = 1
    d_x = d_y = res = 0
    k = 1
    while y > acc:
        d_y += 1
        if x > acc:
            d_x = d_y
        k *= K
        acc += k
    while d_y > d_x:
        y = (y-2)//K+1
        d_y -= 1
        res += 1
    while x != y:
        x = (x-2)//K+1
        y = (y-2)//K+1
        res += 2
    print(res)