트리에서의 다이나믹 프로그래밍 2편 [끝]

안녕하세요, 이번 게시물에서는 트리에서의 다이나믹 프로그래밍 특정 유형을 풀이하는 방법을 다뤄보려고 해요.

 

이 게시물에서 다루는 풀이를 적용하기 위해서는 문제 상황이 아래와 같아야 해요.

1. n개의 정점과 n - 1개의 간선으로 이루어진 트리이어야 한다.

2. n2개의 모든 정점의 쌍인 (i, j)에 대하여 두 개의 정점 (i, j) 최단 거리의 총합을 O(n)에 구해야 한다.

(최단 거리는 꼭 가중치의 합이 아니어도 된다. 가중치의 곱이 될 수도 있다.)

 

위 조건을 만족시키는 문제 네 개 정도에 적용해보았는데, 모두 적용할 수 있었어요.

 

BOJ 1289 - 트리의 가중치 Platinum 3

BOJ 20188 - 등산 마니아 Platinum 5

BOJ 7812 - 중앙 트리 Platinum 5

BOJ 15669 - 트리 위의 입자 Platinum 5

 

이 게시물에서는 BOJ 7812 - 중앙 트리와 BOJ 1289 - 트리의 가중치를 풀이해볼게요.

 

1. BOJ 7812 - 중앙 트리 Platinum 5

이 문제는 아래와 같은 그래프에서 모든 정점과 중앙 정점 사이의 거리의 총합을 구하는 문제예요.

중앙 정점은 모든 정점으로 이르는 비용의 합이 가장 작은 정점이에요.

 

이 문제 상황은 제가 앞서 말씀드린 두 가지 조건을 모두 만족시키죠?

그러면 제 풀이를 적용할 수 있어요.

 

DP 배열은 

dp[i] : i번 정점이 중앙 정점이라고 가정할 때, 모든 정점으로 이르는 비용의 합

와 같이 정의했어요.

따라서 답은 가장 적은 비용의 합을 갖는 min(dp)가 되겠죠?

 

저는 이런 유형의 문제를 두 개의 배열을 이용해서 풀이했어요.

1. subtree_node[i] : i가 루트 노드인 서브 트리의 정점의 개수를 세는 배열 

그림이 조금 이상하지만... 특정 정점을 루트 노드라고 가정하고, 하위 노드의 개수를 저장하는 배열이에요.

subtree_node[i] - 1의 값은 곧 i가 루트 노드인 서브트리의 모든 정점과 i를 잇는 경로의 개수와도 같겠죠?

 

2. child_value[i] :  i가 루트 노드인 서브트리의 모든 정점과 i를 잇는 비용의 총합

위 두 배열은 재귀를 이용해서 어렵지 않게 구할 수 있어요.

중요한 점은 이 두 배열을 이용해서 어떻게 DP 배열을 채울 것인가죠.

 

[풀이]

현재 정점을 i 라고 할 때, 아래 그림과 같이 i의 서브 트리와 서브 트리가 아닌 경우로 나누어서 계산해볼게요.

첫 번째로 현재 정점 now와 서브 트리의 모든 정점까지의 비용 총합은 child_value[now]입니다.

따라서 현재 DP 배열의 값은 dp[now] = child_value[now]입니다.

 

여기까지는 간단하지만, 서브 트리가 아닌 경우는 조금 더 생각을 해보아야 합니다.

 

현재 정점 now와 now의 부모 정점 prev의 관계를 활용해서 서브 트리가 아닌 경우에 dp[now]를 업데이트 할 수 있습니다.

초록색으로 표시된 경로는 prev 정점의 서브 트리가 아닌 경우의 경로입니다.

잘 관찰해보면 now의 서브 트리가 아닌 경우의 경로는 now와 prev를 잇는 간선만 고려하면 구할 수 있다는 것을 알 수 있습니다.

 

now와 prev를 잇는 간선 사이의 비용을 weight라고 하면, dp[now]에

dp[prev] + (subtree_node[0] - subtree_node[now] - 1) * weight

을 더해주면 됩니다.

 

subtree_node[0] - subtree_node[now] - 1의 값은

(0번 정점을 루트로 하는 서브 트리의 정점의 개수(전체 트리)) - (now번 정점을 루트로 하는 서브 트리의 정점의 개수) -1

이고, 이는 곧 초록색 경로의 개수라고 이해할 수 있습니다.

 

여기에 비용을 곱해 더하면 now에서 서브 트리가 아닌 정점으로 가는 비용의 총합을 구할 수 있습니다.

 

하지만 여기서 추가로 고려해야 할 것이 있습니다.

DP 배열은 해당 정점으로부터 모든 정점으로 이르는 비용의 총합이기 때문에 파란색 경로도 포함되어 있습니다.

이 경우는 위에서 dp[now] = child_value[now]와 같이 설정해주면서 이미 고려했으므로 파란색 화살표의 경우는 빼주어야 합니다.

 

이 경우는 child_value[now] + subtree_node[now] * weight로 계산해줄 수 있습니다.

수식이 다소 복잡하지만 그림으로 그려보면 이해가 될 거라고 생각해요.

 

따라서 서브 트리가 아닌 경우의 DP 배열의 값은

dp[now] += dp[prev] + (subtree_node[0] - subtree_node[now] - 1) * weight - (child_value[now] + subtree_node[now] * weight)

입니다.

 

이런 방식으로 현재 정점에 대한 업데이트를 마친 후, 재귀적으로 모든 정점에 대해 업데이트 해주면 됩니다.

위 과정이 끝난 후, min(dp)를 출력하면 정답을 맞을 수 있습니다.

import sys
sys.setrecursionlimit(int(1e6)) 

while True:
    n = int(input())
    if n == 0: break

    g = [[] for _ in range(n)]
    for _ in range(n - 1):
        u, v, w = map(int, input().split())
        g[u].append([v, w])
        g[v].append([u, w])
    
    vis = [False for _ in range(n)]
    c = [0 for _ in range(n)]
    r = [0 for _ in range(n)]
    
    def sub(now):
        vis[now] = True
        for _ in g[now]:
            nxt, w = _[0], _[1]
            if vis[nxt]: continue
            temp = sub(nxt)
            r[now] += r[nxt] + 1
            c[now] += temp + r[nxt] * w + w
        return c[now]
    sub(0)
    vis = [False for _ in range(n)]
    
    dp = [0 for _ in range(n)]

    def dfs(now, prev, weight):
        global ans
        vis[now] = True
        
        dp[now] = c[now]
        if prev != -1:
            dp[now] += dp[prev] + (r[0] - r[now] - 1) * weight - (c[now] + r[now] * weight)
        for _ in g[now]:
            nxt, w = _[0], _[1]
            if vis[nxt]: continue
            dfs(nxt, now, w)
        
    dfs(0, -1, -1)
    print(min(dp))

 

2. BOJ 7812 - 트리의 가중치 Platinum 3

저번 문제는 가중치를 거리라고 했으면, 이번 문제는 경로에 해당하는 간선의 곱으로 정의했습니다.

모든 n2개의 순서쌍에 대하여 가중치의 총합을 구하는 문제입니다.

 

이 문제 역시 제가 말씀드린 문제의 조건에 부합한다고 할 수 있습니다.

 

이번엔 DP 배열을

dp[i] = i번 정점을 거쳐가는 모든 경우의 가중치의 총합

이라고 정의해보겠습니다.

 

i번 정점을 거쳐가는 경우를 생각해보면, i번 정점은 무조건 서브 트리의 루트 노드가 되어야 한다는 것을 알 수 있습니다.

따라서 난이도는 Platinum 3으로 높아졌지만 오히려 저번 문제처럼 부모 노드를 생각하지 않아도 돼서 풀기 더 편하다고 할 수 있습니다.

 

그래서 이번에는 한 개의 배열 만을 활용해서 풀어보겠습니다.

child_value[i] :  i가 루트 노드인 서브트리의 모든 정점과 i를 잇는 가중치의 총합

 

위 문제와 비슷한 컨셉이니 이 배열을 채우는 것은 그리 어렵지 않게 할 수 있을 것이라고 생각합니다.

[풀이]

이번에는 now와 next의 관계를 통해 next가 루트인 서브 트리의 모든 정점에서 now에 도달하는 가중치의 총합을 구해봅시다.

 

child_value[next]를 이루는 k개의 경로에 대하여 각 가중치를 w1, w2, ..., wk라고 할 때, now에 도달하는 가중치는 w1 * w + w2 * w + ... + wk * w + w입니다. 이는 곧 child_value[next] * w + w로 정리할 수 있습니다.

따라서

(next가 루트인 서브 트리의 모든 정점에서 now에 도달하는 가중치의 총합): child_value[next] * w + w

입니다.

 

여기서 추가적인 변수 temp를 선언해줍시다.

temp는 now의 자식 정점 중 next 정점을 살펴 보고 있을 때, 지금까지 살펴본 모든 정점에 대하여 해당 정점을 루트 노드로 가지는 서브 트리의 모든 정점에서 now에 도달하는 가중치의 총합이라고 정의해봅시다.

 

말로 설명하니 어려운데, 아래 그림에서 왼쪽에서 오른쪽 순서로 살펴본다고 할 때, 빨간색으로 되어있는 부분이라고 생각하시면 돼요.

 

위에서 계산했던 (next가 루트인 서브 트리의 모든 정점에서 now에 도달하는 가중치의 총합): child_value[next] * w + w을 s라고 가정하면 아래와 같은 과정으로 next 정점까지 고려하여 dp 배열을 업데이트 할 수 있습니다.

따라서 모든 자식 정점에 대하여 dp[now] += temp * s + s를 해주면 정답을 구할 수 있습니다.

(모듈러 연산은 생략했습니다.)

DP 배열을 업데이트 한 뒤, 다음 정점을 살펴볼 것을 고려하여 temp += s로 temp 변수도 업데이트 해주어야 합니다.

 

위 과정을 재귀적으로 모든 노드에 대해 수행해주면 DP 배열을 모두 채울 수 있습니다.

모든 과정이 끝난 뒤, 모든 정점 순서쌍에 대한 가중치의 총합을 구해야 하므로 sum(dp)가 정답이 됩니다.

 

import sys
sys.setrecursionlimit(int(1e6)) 

n = int(input())
g = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    u, v, w = map(int, input().split())
    g[u].append([v, w])
    g[v].append([u, w])

MOD = int(1e9) + 7

vis = [False for _ in range(n + 1)]
c = [0 for _ in range(n + 1)]
ans = 0

def sub(now):
    vis[now] = True
    for _ in g[now]:
        nxt, w = _[0], _[1]
        if vis[nxt]: continue
        c[now] = (c[now] + sub(nxt) * w + w) % MOD
    return c[now]
    
sub(1)
vis = [False for _ in range(n + 1)]

dp = [0 for _ in range(n + 1)]

def dfs(now):
    global ans
    vis[now] = True
    
    temp = 0
    for _ in g[now]:
        nxt, w = _[0], _[1]
        if vis[nxt]: continue
        s = c[nxt] * w + w
        dp[now] = (dp[now] + temp * s + s) % MOD
        temp += s
        dfs(nxt)

dfs(1)
print(sum(dp) % MOD)

---

트리에서의 다이나믹 프로그래밍은 아주 어려운 유형입니다.

그래서 최대한 일관적인 방법으로 문제에 접근해보고 싶었습니다.

 

만약 문제가 제가 설명드린 이 유형의 조건에 부합한다면 서브 트리와 서브 트리가 아닌 경우의 상태를 담고 있는 배열을 선언하여, 이들의 조합으로 정답을 구할 수 있습니다.

이 유형 풀이의 시간 복잡도는 O(n)입니다.

 

물론 더 쉽고 획기적인 풀이가 있을 수 있지만 같은 유형에 해당하는 일관적인 풀이에 초점을 둔 것을 알려드립니다.

---

제 글이 많은 도움이 되었길 바라며,
긴 글 봐주셔서 감사합니다!
멋진 개발자가 되기 위해 더 열심히 달리겠습니다!
- 달맹 -