성대나라의 물탱크

• 복잡한 듯 보이지만, 각 노드는 자식 노드에 물이 채워질때마다 업데이트되며, 그때 업데이트 되는 값은 그 노드의 깊이(=루트로부터의 거리)이다.
• 결국 특정 노드의 채워진 물의 양은, {자식 노드들에 쿼리가 들어온 횟수} * {그 노드의 깊이} 가 된다.
• 깊이를 곱하는 것만 제외하면 회사 문화 3과 같은 문제이고, 똑같이 오일러 경로 테크닉과 구간합 쿼리를 조합해서 풀 수 있다.
• 회사 문화 3에 비해 추가로 필요한, 각 노드의 깊이를 구하는 것은, 오일러 경로 테크닉에서 트리를 DFS traverse할때 함께 계산해줄 수 있고, 추가적인 시간 복잡도가 들지는 않는다.
• 그래서 시간 복잡도도 회사 문화 3와 똑같이 O(N+MlogN)이 된다.

"""Solution code for "BOJ 18227. 성대나라의 물탱크".

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

import sys
from teflib import fenwicktree
from teflib import tgraph


def main():
    N, C = [int(x) for x in sys.stdin.readline().split()]
    graph = [[] for _ in range(N)]
    for _ in range(N - 1):
        x, y = [int(x) for x in sys.stdin.readline().split()]
        graph[x - 1].append(y - 1)
        graph[y - 1].append(x - 1)

    order = level = 0
    subtree_ranges = [[None, None] for _ in range(N)]
    levels = [None] * N
    for node in tgraph.dfs(graph, C - 1, yields_on_leave=True):
        if subtree_ranges[node][0] is None:
            subtree_ranges[node][0] = order
            order += 1
            level += 1
            levels[node] = level
        else:
            subtree_ranges[node][1] = order
            level -= 1
    fenwick = fenwicktree.FenwickTree(N)

    Q = int(sys.stdin.readline())
    for _ in range(Q):
        query_type, A = [int(x) for x in sys.stdin.readline().split()]
        if query_type == 1:
            fenwick.update(subtree_ranges[A - 1][0], 1)
        else:
            print(fenwick.query(*subtree_ranges[A - 1]) * levels[A - 1])


if __name__ == '__main__':
    main()

• Dependency
  • teflib.fenwicktree.FenwickTree
  • teflib.tgraph.dfs