회사 문화 4

• 회사 문화 2와 회사 문화 3을 합쳐 놓은 문제
• 부하방향 칭찬과 상사방향 칭찬을 동시에 계산 가능하도록 하려고 고심할 필요가 없다. 그냥 부하방향으로 받은 칭찬과 상사방향으로 받은 칭찬을 따로따로 처리한 뒤에, 칭찬값의 합을 출력할때에는 두 가지를 더해서 출력하면 된다.
• 즉, 오일러 경로 테크닉을 적용한 뒤에, 부하방향으로 받는 칭찬은 회사 문화 2에서 설명했듯이 구간합 업데이트+포인트 쿼리로, 상사방향으로 받는 칭찬은 회사 문화 3에서 설명했듯이 포인트 업데이트+구간합 쿼리로 처리하면 된다. 두가지 용도의 펜윅트리를 따로 만들어서 쓰면 된다.
• 부하방향 칭찬이든 상사방향 칭찬이든, 업데이트와 쿼리 모두 O(logn)이 걸린다. 총 시간 복잡도는 오일러 경로 테크닉에 O(n), m개의 쿼리 처리에 O(mlogn), 합쳐서 O(n+mlogn)이다

"""Solution code for "BOJ 14288. 회사 문화 4".

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

import sys
from teflib import fenwicktree
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, m = [int(x) for x in sys.stdin.readline().split()]
    tree = [[] for _ in range(n)]
    bosses = [int(x) for x in sys.stdin.readline().split()]
    for employee, boss in enumerate(bosses):
        if boss != -1:
            tree[boss - 1].append(employee)
    subtree_ranges = euler_tour_technique(tree, 0)
    fenwick_for_upward = fenwicktree.FenwickTree(n + 1)
    fenwick_for_downward = fenwicktree.FenwickTree(n + 1)
    is_downward = True
    for _ in range(m):
        query = [int(x) for x in sys.stdin.readline().split()]
        if query[0] == 1:
            _, i, w = query
            beg, end = subtree_ranges[i - 1]
            if is_downward:
                fenwick_for_downward.update(beg, w)
                fenwick_for_downward.update(end, -w)
            else:
                fenwick_for_upward.update(beg, w)
        elif query[0] == 2:
            i = query[1]
            beg, end = subtree_ranges[i - 1]
            print(fenwick_for_downward.query(0, beg + 1)
                  + fenwick_for_upward.query(beg, end))
        else:
            is_downward = not is_downward


if __name__ == '__main__':
    main()

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