Tree With One Edge

• 트리에서 돌을 이동시키는 게임의 변형이다. 트리에 에지를 하나 더 추가했을때 승패가 어떻게 될지를 생각해보면 된다.
• 노드 u의 승리포지션/패배포지션이 바뀌도록 에지를 추가하는 방법의 수를 dp[u]로 정의해서 tree dp로 구할수 있다.
• u의 자식노드들을 v1, v2, … 이라고 하자.
• 자식노드중에 패배포지션이 1개라면,
  • u는 승리포지션이다.
  • u를 패배포지션으로 바꾸려면
    • 패배포지션인 자식노드가 승리포지션이 되도록 바꿔주면 된다. 그러면 자식노드중에 패배포지션이 하나도 없어지므로 u가 패배포지션이 된다
    • dp[u] = dp[v_i]. (v_i는 패배포지션)
• 자식노드중에 패배포지션이 2개 이상이라면,
  • u는 승리포지션이다.
  • u를 패배포지션으로 바꾸는 것은 불가능하다.
    • 어떤 자식노드 하나를 바꿔도, 자식노드중에 패배포지션이 한개 이상 있으므로 u는 여전히 승리포지션이다.
    • dp[u] = 0
• 자식노드중에 패배포지션이 하나도 없으면:
  • u는 패배포지션이다.
  • u를 승리포지션으로 바꾸려면,
    • u에서 패배포지션인 노드로 가는 에지를 추가하거나, 자식노드 하나를 승리포지션으로 바꾸면 된다.
    • dp[u] = C + (dp[v1] + dp[v2] + …)
• 이렇게 dp를 통해서 루트노드의 승패를 바꿀수 있는 에지의 개수를 구할수 있고, 승패를 바꾸지 못하는 에지의 개수 또한 전체 가능한 에지의 가짓수 n^2에서 저 값을 빼줌으로써 구할수 있다. 루트노드가 승리포지션이면, 승패를 바꾸지 못하는 에지의 개수가 답이 되고, 루트노드가 패배포지션이면 승패를 바꿀수 있는 에지의 개수가 답이 된다.
• 총 시간복잡도는 O(n)

"""Solution code for "BOJ 32744. Tree With One Edge".

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

Tags: [game theory]
"""

import sys
from teflib import psutils
from teflib import tree as ttree

ROOT = 0


def solve_dp(dp_children, _):
    loss_pos_count = sum(1 for is_win_pos, _ in dp_children if not is_win_pos)
    if loss_pos_count > 0:
        is_win_pos_cur = True
        if loss_pos_count > 1:
            change_pos_count = 0
        else:
            change_pos_count = next(
                count for is_win_pos, count in dp_children if not is_win_pos
            )
    else:
        is_win_pos_cur = False
        change_pos_count = sum(count for _, count in dp_children) + 1
    return (is_win_pos_cur, change_pos_count)


@psutils.run_n_times
def main():
    n = int(sys.stdin.readline())  # pylint: disable=unused-variable
    parents = [0] + [int(x) - 1 for x in sys.stdin.readline().split()]

    dp = ttree.dp_on_rooted_tree(ttree.RootedTree(ROOT, parents), solve_dp)
    loss_node_count = sum(1 for is_win_pos, _ in dp if not is_win_pos)
    is_win_pos, change_pos_count = dp[ROOT]
    change_edge_count = change_pos_count * loss_node_count
    if is_win_pos:
        print(n * n - change_edge_count)
    else:
        print(change_edge_count)


if __name__ == '__main__':
    main()

• Dependency: teflib.tree.dp_on_tree