트리 색칠하기

• 단순하게 생각하면, 1번과 2번 색깔로만 트리를 색칠하면 될것 같지만 그렇게는 안된다. 반례는 A노드와 B노드가 연결되어 있고, A노드와 B노드에는 각각 10개씩의 리프노드가 붙어있는 트리를 떠올려보면, 리프노드를 모두 1로 색칠하고 A와 B노드를 2와 3으로 색칠해야 한다는 것을 알 수 있다.
• 그렇다면, DP를 써서 푸는 방법이 떠오른다. 각 서브트리에 대해서 루트를 k로 색칠했을때의 최소 코스트를 각각 계산해놓고, 그것들로부터 상위 트리의 값을 계산하는 방법이다. ]
• 하지만 이 방법은, 가능한 색칠 방법의 수가 n개이기 때문에, 하나의 노드에 대해서 n개의 값을 계산해야 한다. 그렇게 하면 총 시간복잡도가 O(n^2)이 되어서 시간 내에 풀 수 없다
• 관찰을 통해 알아낼 수 있는 것은, n개의 노드가 있는 트리를 최소코스트로 색칠할때에 필요한 색깔의 수는 logn개를 넘지 않는다는 것이다. (증명은 여기). 따라서 색의 수를 logn으로 제한하면, 총 시간 복잡도를 O(nlogn)으로 줄일수 있다.
  • 이것이 거의 정해처럼(?) 받아들여지는 풀이인것 같다. 난이도 기여에 적힌 커멘트들도 전부 이 풀이를 기준으로 했다.
• 그러나 저런 관찰과 증명 없이도, O(n) 풀이를 만들수 있다. 이 풀이는 다른 블로그에서는 본적이 없다..
• 방법은 각 서브 트리에 대해서 최소 비용(cost[i]), 최소 비용일때의 루트의 색깔(color[i]), 두번째로 작은 비용(cost2[i]) 의 세가지 값을 저장하는 것이다. 서브트리들의 이 값을 가지고 상위 트리의 값을 계산할때는, color[i] 에 포함되지 않은 색깔중 가장 작은 색으로 루트를 색칠하는 경우와, color[k]로 루트를 칠하는 경우를 구해주면 된다.
  • 전자는 그냥 sum(cost[i]) + {color[i] 에 포함되지 않은 가장 작은 색} 으로 계산하면 된다
  • 후자는 color[i] 와 color[k]가 같으면 cost2[i]를, 다르면 cost[i] 를 모든 i에 대해서 더하고 루트에 칠하는 color[k]값을 더해주면 된다. sum(cost[i])에서 (cost2[k]-cost[k]) 들을 더해주는 식으로 구현하면 O(n)에 계산이 가능하다.
  • x개의 차일드 노드들로부터 부모 노드의 DP값을 구하는 것을 O(x)에 할수 있으므로, 전체 트리에 대한 계산은 O(n)에 가능하다.

"""Solution code for "BOJ 1693. 트리 색칠하기".

- Problem link: https://www.acmicpc.net/problem/1693
- Solution link: http://www.teferi.net/ps/problems/boj/1693

Tags: [DP] [DP on Tree]
"""

import collections
import itertools
import sys
from teflib import ttree


def calc_dp(subtree_vals, _):
    increments = collections.defaultdict(int)
    sub_cost_sum = 0
    for sub_min_cost, sub_min_color, sub_second_min_cost in subtree_vals:
        sub_cost_sum += sub_min_cost
        increments[sub_min_color] += sub_second_min_cost - sub_min_cost
    unused_colors = (x for x in itertools.count(1) if x not in increments)
    min_inc = min_color = next(unused_colors)
    second_min_inc = next(unused_colors)

    for color, inc in increments.items():
        new_inc = inc + color
        if new_inc < min_inc:
            min_inc, min_color, second_min_inc = new_inc, color, min_inc
        elif new_inc < second_min_inc:
            second_min_inc = new_inc
    min_cost = sub_cost_sum + min_inc
    second_min_cost = sub_cost_sum + second_min_inc
    return (min_cost, min_color, second_min_cost)


def main():
    n = int(sys.stdin.readline())
    tree = [[] for _ in range(n)]
    for _ in range(n - 1):
        u, v = [int(x) for x in sys.stdin.readline().split()]
        tree[u - 1].append(v - 1)
        tree[v - 1].append(u - 1)

    print(ttree.dp_on_tree(tree, calc_dp, 0)[0])


if __name__ == '__main__':
    main()

• Dependency: teflib.ttree.tree_on_dp