[BOJ 15899]트리와 색깔

[BOJ 15899]트리와 색깔 - P2

문제 설명

노드가 N개인 트리가 주어진다. 각 노드는 각자 i번의 색깔로 칠해져있다. 그리고 한 가지 쿼리가 주어진다. 쿼리는 a번 노드의 subtree에서 색칠된 색깔의 번호가 b이하인 노드의 개수가 몇 개인지 묻고있다. 모든 쿼리의 정답을 다 더해서 모듈러를 취한 뒤 답하면 된다.

문제 풀이

트리에서 특정 값 x보다 작은 노드의 개수? 이거는 1초의 고민할 필요도 없이 머지소트 트리가 떠올라야 한다.
맨 처음에 제출할 때는 쌩으로 머지소트 트리를 만들었다. 쌩으로 머지소트 트리를 만들면 linear한 트리가 입력으로 주어질 때, 머지소트 트리가 저장해야할 칸이 $O(N^2)$이므로 메모리가 터지게 된다. 따라서 머지소트 트리를 만들기 전에 이 문제를 해결할 전처리를 해야 한다.
이 문제도 '자동차 공장'문제 처럼 i번 노드의 자식들에 대한 쿼리를 log만에 처리해야 하므로 오일러 순회 기법(ETT)를 한다.
ETT를 거친 후에 새롭게 배정된 노드 번호를 바탕으로 머지소트 트리를 구성하고, 쿼리에 답하면 된다. 쿼리에 답할 때는 ETT를 돌린 후에 얻는 정보인 'i번 노드의 자식 중에서 가장 작은 인덱스와 큰 인덱스'를 바탕으로 getSum(l[x],r[x])로 물어보면 된다.

소스 코드

#include<stdio.h>
#include<algorithm>
#include<vector>
using namespace std;
typedef long long ll;

const int SIZE = 200009;
const int MOD = 1e9 + 7;
vector<int> tree[SIZE*8],graph[SIZE];
int startIdx=1,cnt = 0;
int l[SIZE],r[SIZE],c[SIZE];

void makeMergeSortTree(int node){
    if(node >= startIdx) return; // leaf node

    makeMergeSortTree(node*2);
    makeMergeSortTree(node*2+1);

    for(int i : tree[node*2])
        tree[node].push_back(i);

    for(int i : tree[node*2+1])
        tree[node].push_back(i);

    sort(tree[node].begin(),tree[node].end());

    return;
}

int query(int node,int cl,int cr,int ql,int qr,int x){
        if(cl > qr || ql > cr) return 0;
        if(qr >= cr && ql <= cl){
            int cnt = (upper_bound(tree[node].begin(), tree[node].end(),x) - tree[node].begin());
            return cnt;
        }
        int mid = (cl+cr)/2;
        return query(node*2,cl,mid,ql,qr,x) + query(node*2+1,mid+1,cr,ql,qr,x);
}

void ett(int node,int p){
    l[node] = ++cnt;
    for(int next : graph[node])
        if(next != p) ett(next,node);
    r[node] = cnt;
}

int main(){
    int n,q,k;    scanf("%d %d %d",&n,&q,&k);

    for(; startIdx < n ; startIdx *= 2);
    for(int i=1 ; i<=n ; i++){
        scanf("%d",&c[i]);
    }

    for(int i=0 ; i<n-1 ; i++){
        int a,b;    scanf("%d %d",&a,&b);
        graph[a].push_back(b);
        graph[b].push_back(a);
    }

    ett(1,-1);
    for(int i=1 ; i<=n ; i++){
      tree[startIdx + l[i] - 1].push_back(c[i]);
    }
    for(int i=n ; i<=startIdx ; i++){
        tree[startIdx + i].push_back(2e9);
    }
    makeMergeSortTree(1);

    ll sum = 0;
    while(q--){
        int v,c;    scanf("%d %d",&v,&c);
        sum += query(1,1,startIdx,l[v],r[v],c);
    }

    printf("%lld\n",sum % MOD);
    return 0;
}