트리와 색깔

• 오일러 경로 테크닉과 구간 Rank 쿼리를 결합한 문제.
• 기본적인 풀이 방식을 결합해서 풀면 끝이다. 오일러 경로 테크닉에 O(N). 구간 Rank 쿼리를 오프라인 쿼리로 처리해서 푸는 데에 O((N+M)log(C)). 합치면 O((N+M)log(C)).
• 다만 구현을 조금 컴팩트하게 할수 있는 방법이 있다. 구현하는 과정에서 앞서말한 풀이 방식을 그냥 순서대로 결합하면, 오일러 경로 테크닉을 사용해서 노드들을 DFS order로 재배치하고, 각 서브트리에 대응되는 구간들을 모두 구하고, 실제 서브트리 쿼리들을 구간 쿼리로 변환해서 저장하고, dfs order로 저장된 노드들을 order statistic tree에 추가하면서 쿼리들을 처리하는 방식이 된다. 하지만, 조금 손을 보면, 서브트리 쿼리들을 모아서 저장한 뒤에, DFS로 노드들을 방문하면서 곧바로 order statistic tree에 추가하고, 서브트리 쿼리들을 처리해 나갈 수도 있다. 속도는 별 차이 없지만, 코드가 좀더 짧아진다.
• 이 문제는 제한시간은 2초지만, python에 추가 시간이 주어지지 않는다.. Python으로는 2초 안에 풀리지 않았고, PyPy로 통과했다

"""Solution code for "BOJ 15899. 트리와 색깔".

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

This code should be submitted with PyPy3. Python3 gets TLE.
"""

import sys
from teflib import fenwicktree
from teflib import tgraph

MOD = 1_000_000_007


def main():
    N, M, C = [int(x) for x in sys.stdin.readline().split()]
    colors = [int(x) for x in sys.stdin.readline().split()]
    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)
    queries = [[] for _ in range(N)]
    for _ in range(M):
        v, c = [int(x) for x in sys.stdin.readline().split()]
        queries[v - 1].append(c)
    
    answer = 0
    color_set = fenwicktree.OrderStatisticTree(C)
    is_visited = [False] * N
    for node in tgraph.dfs(tree, 0, is_visited=is_visited, 
                           yields_on_leave=True):             
        if not is_visited[node]:
            for c in queries[node]:
                answer -= color_set.count_less_than(c + 1)            
            color_set.add(colors[node])
        else:
            for c in queries[node]:
                answer += color_set.count_less_than(c + 1)    

    print(answer % MOD)


if __name__ == '__main__':
    main()

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