회사 문화 5

• 오일러 경로 테크닉과 구간 업데이트+구간 쿼리를 결합한 문제.
• 우선 오일러 경로 테크닉을 사용해서, 서브트리에 대한 쿼리를 구간에 대한 쿼리로 변환한다. 구간에 대해서 처리해야 하는 쿼리의 종류는 assignment 업데이트 + add 쿼리로, lazy propagation을 지원하는 세그먼트 트리로 쉽게 처리 가능하다.
• 오일러 경로 테크닉에 O(n), 세그먼트 트리 구축에 O(n), m개의 쿼리를 각각 O(logn)에 처리하는 데에 O(mlogn). 총 O(n+mlogn)의 시간 복잡도로 풀 수 있다.
• 실수하기 쉬운 부분.
  • 이전 회사 문화 문제들(회사 문화 2, 회사 문화 3, 회사 문화 3) 과는 달리 대표의 상사가 -1이 아닌 0으로 주어진다. 신경 안쓰로 기존 코드 잘못 복붙하다가는 이것때문에 틀릴 수 있다
  • 모든 쿼리는 상사 본인은 제외하고 부하들에 대해서만 적용된다. 구간에서 서브트리루트에 해당하는 값을 제외하고 처리해야 한다.
  • 부하가 하나도 없는 사람에게도 쿼리가 들어올 수 있다.

"""Solution code for "BOJ 18437. 회사 문화 5".

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

import operator
import sys
from teflib import segmenttree
from teflib import tgraph


def euler_tour_technique(tree, root):
    subtree_ranges = [[None, None] for _ in tree]
    order = 0
    for node in tgraph.dfs(tree, root, yields_on_leave=True):
        if subtree_ranges[node][0] is None:
            subtree_ranges[node][0] = order
            order += 1
        else:
            subtree_ranges[node][1] = order
    return subtree_ranges


def main():
    N = int(sys.stdin.readline())
    tree = [[] for _ in range(N)]
    bosses = [int(x) for x in sys.stdin.readline().split()]
    for employee, boss in enumerate(bosses):
        if boss != 0:
            tree[boss - 1].append(employee)
    subtree_ranges = euler_tour_technique(tree, 0)
    lazy_segtree = segmenttree.LazySegmentTree(
        [1] * N,
        merge=operator.add,
        update_value=lambda v, p, size: size if p else 0,
        update_param=lambda p1, p2: p2)

    M = int(sys.stdin.readline())
    for _ in range(M):
        query_type, i = [int(x) for x in sys.stdin.readline().split()]
        beg, end = subtree_ranges[i - 1]
        if beg + 1 == end:
            if query_type == 3:
                print(0)
            continue
        if query_type == 1:
            lazy_segtree.range_update(beg + 1, end, True)
        elif query_type == 2:
            lazy_segtree.range_update(beg + 1, end, False)
        else:
            print(lazy_segtree.query(beg + 1, end))


if __name__ == '__main__':
    main()

• Dependency
  • teflib.segmenttree.LazySegmentTree
  • teflib.tgraph.dfs