[Python/파이썬] 백준 14942번: 개미

문제

 

풀이

1. 주어진 입력은 양방향 간선이므로 단방향 간선으로 만들어줘야함

2. 희소배열 사용

https://hello-capo.netlify.app/algorithm-sparse-table/

🧩 [Algorithm] 희소 배열(Sparse Table) 알고리즘 | hello-capo!

3. LCA와 유사하게 노드 탐색

https://cyberflower.github.io/2019/07/22/LCA.html

LCA(Least Common Ancestor) - CyberFlower Algorithm

 

코드

import sys, math
input = sys.stdin.readline

#양방향 그래프를 단방향으로 바꿈
def dfs(x):
    visited[x] = True
    for next, power in tmp[x]:
        if not visited[next]:
            arr[next] = [x, power]
            dfs(next)

#희소 행렬 이용
def lca(j, w):
    for i in range(int(math.log2(n)), -1, -1):
        if spareTable[i][j][1] <= w and spareTable[i][j][0] != 0:
            w -= spareTable[i][j][1]
            j = spareTable[i][j][0]
    return j
    
n = int(input().rstrip())
ant = [0] + [int(input().rstrip()) for _ in range(n)] #개미 에너지
arr = [[] for _ in range(n + 1)] #단방향 그래프
tmp = [[] for _ in range(n + 1)] #양방향 그래프(임시)
visited = [False] * (n + 1) #방문 여부

#양방향 그래프 입력
for _ in range(n - 1):
    a, b, c = map(int, input().rstrip().split())
    tmp[b].append([a, c])
    tmp[a].append([b, c])

#단방향 그래프로 만들어줌
dfs(1)
arr[1] = [1, 0]

#희소 행렬 생성
spareTable = [[[0, 0] for _ in range(n + 1)] for _ in range(int(math.log2(n)) + 1)]
for i in range(1, n + 1):
    spareTable[0][i] = arr[i]

for k in range(1, int(math.log2(n)) + 1):
    for i in range(1, n + 1):
        next, power = spareTable[k - 1][i]
        spareTable[k][i] = [spareTable[k - 1][next][0], power + spareTable[k - 1][next][1]]

for i in range(1, n + 1):
    print(lca(i, ant[i]))

 

마무리

개인적으로 너무 어려웠다... 구글링해서 겨우 품