백준 14942번: 개미 (Sparse Table)

문제링크:14942번: 개미 (acmicpc.net)

14942번: 개미

이 문제는 루트가 1인 n개의 노드를 가진 트리가 있을 때, 각 노드에서 무조건 루트 방향(조상노드방향) 으로 이동한다고 했을 때, 그리고 갈 수 있는 거리가 한정되어 있을 때 출발 해서 도달할 수 있는 최대 조상 노드를 구하는 문제이다. 이 문제는 얼마 전에 푼 LCA문제에서 이진 탐색을 이용한 희소 배열을 이용해서 조상 노드로 이동하는 원리를 사용했다.

 

먼저 dfs를 통해 루트(노드0)부터 탐색을 이어가며 parent[n][k][0]은 n노드에서 2**k만큼 올라간 조상 노드를 저장하고 있고, parent[n][k][1]은 그 노드까지 가는데 걸리는 거리를 저장하고 있다. 혹시 이후 탐색에 방해가 될 수 있으니 초기의 parent[n][k][0]은 존재하지 않는 노드인 -1, parent[n][k][1]은 주어진 체력으로 절대 갈 수 없는 100001로 정해주었다.  new노드를 발견하면 먼저 new의 1단계 위 부모는 root이므로 parent[new][0][0]=root, 거리는 newdis이므로 parent[new][0][1]=newdis 를 해주고 그 위는 DP방식을 이용해 2**k만큼 올라간 조상 노드가 -1(끝) 이 될떄까지 채워나갔다.

import sys

N=int(sys.stdin.readline().rstrip())

maxdepth=1
while 2**(maxdepth)<N:
    maxdepth+=1
    
ants=[]
graph=[[] for _ in range(N)]

for _ in range(N):
    ants.append(int(sys.stdin.readline().rstrip()))

for _ in range(N-1):
    a,b,c=map(int,sys.stdin.readline().rstrip().split())
    a-=1
    b-=1
    graph[a].append([b,c])
    graph[b].append([a,c])

visited=[0 for _ in range(N)]
visited[0]=1
parent=[[[-1,100001] for _ in range(maxdepth)] for _ in range(N)]

def dfs(root):
    for new,newdis in graph[root]:
        if visited[new]==1:
            continue
        visited[new]=1
        parent[new][0][0]=root
        parent[new][0][1]=newdis
        for d in range(1,maxdepth):
            parent[new][d][0]=parent[parent[new][d-1][0]][d-1][0]
            if parent[new][d][0]==-1:
                break
            parent[new][d][1]=parent[new][d-1][1]+parent[parent[new][d-1][0]][d-1][1]

        dfs(new)
    return

dfs(0)

def find_room(cur):
    energy=ants[cur]
    if cur==0:
        return cur
    
    for d in range(maxdepth-1,-1,-1):
        new=parent[cur][d][0]
        move=parent[cur][d][1]
        if new==-1:
            continue
        if move>energy:
            continue
        energy-=move
        cur=new
        if cur==0:
            return cur
    return cur

for i in range(N):
    ans=find_room(i)
    print(ans+1)

find_room()함수를 통해 최대 도달할 수 있는 노드를 반환했다. maxdepth-1부터 0까지 조사하며 만약 new가 -1이라면 루트 노드를 지나쳤다는 뜻이기에 제외해주고, 만약 move가 현재 가진 energy보다 크다면 갈 수 없으므로 제외해준다. 그렇게해서 갈 수 있는 노드가 나오면 energy에서 move만큼 빼주고, cur을 new로 업데이트 해 가며 만약 루트 노드 0에 다다르면 반환, 아니라면 for문을 마치고 현재 cur을 반환한다.

 

처음에 move를 구하기 위해 사용한 방법은 dislst를 만들어 처음 dfs()함수를 하며 노드를 지날 때 거리를 누적 해 가며 0부터 각 노드까지의 걸리는 총 거리를 저장해두었다. 그러고 find_room()에서 탐색 할 때 dislst[cur]-dislst[new]를 하는 방식으로 move를 구했었지만 테스트 케이스와 내가 생각해 본 케이스들은 모두 맞았지만 정작 틀렸습니다를 받았다. move 계산하는 방법만 parent리스트에 같이 저장하는 방법으로 바꿨더니 정답을 받을 수 있었다.

 

생각해 보니 굴의 길이는 최대 10,000이고, 굴의 개수는 최대 100,000개이므로 총 길이는 1,000,000,000 로 매우 큰 값이 나왔을 것이고 큰 수를 계산하는데 시간이 낭비되거나 계산이 제대로 되지 않았나 보다.