[BOJ 13518] 트리와 쿼리 9

풀이는

http://weeklyps.com/entry/%ED%8A%B8%EB%A6%AC%EC%97%90%EC%84%9C%EC%9D%98-%EB%AA%A8%EC%8A%A4-%EC%95%8C%EA%B3%A0%EB%A6%AC%EC%A6%98-%ED%99%9C%EC%9A%A9-Mos-algorithm-on-tree

이 곳에 서술 되어 있으므로, 상세한 사항은 아래의 소스로 대신한다.

// code by RiKang, weeklyps.com
#include<stdio.h>
#include<algorithm>
#include<vector>
#define PII pair<int,int>
#define PIII pair<PII,int>
using namespace::std;
 
const int N=300005, M = 1000005, SQRT_N = 400;
int n;
int a[N];
vector<int> adj[N];
bool visit[N];
 
int arr_n;        // arr의 길이
int arr[N];  // 트리를 펼친 배열
int st[N];   // st[x] = (arr[i] == x)인 최소 i
int en[N];   // en[x] = (arr[i] == x)인 최대 i
 
PII query[N];
PIII order[N];
int cnt[M];
int cnt2[M];
int q;
int an;
int ans[M];
int lc[M];
 
void add2(int idx){ // idx 가 구간에서 1번 등장하게 됨.
    if(cnt2[a[idx]]==0) an++;
    cnt2[a[idx]]++;
}
 
void erase2(int idx){ // idx 가 구간에서 짝수번 등장하게 됨.
    cnt2[a[idx]]--;
    if(cnt2[a[idx]]==0) an--;
}
 
void add(int idx){
    if(cnt[arr[idx]]==0) add2(arr[idx]); 
    if(cnt[arr[idx]]==1) erase2(arr[idx]);
    cnt[arr[idx]]++;
}
 
void erase(int idx){
    cnt[arr[idx]]--;
    if(cnt[arr[idx]]==1) add2(arr[idx]);
    if(cnt[arr[idx]]==0) erase2(arr[idx]); 
}
 
void go_query(PII prev, PII now){
    int l1 = prev.first;
    int r1 = prev.second;
    int l2 = now.first;
    int r2 = now.second;
    for(int i = l1-1; i>=l2; i--) add(i); // l1 > l2 인 경우, 구간이 늘어났으므로 노드 추가
    for(int i = r1+1; i<=r2; i++) add(i); // r1 < l2 인 경우, 구간이 늘어났으므로 노드 추가
    for(int i = l1; i<l2; i++) erase(i); // l1 < l2 인 경우, 구간이 줄었으므로 노드 삭제
    for(int i = r1; i>r2; i--) erase(i); // r1 < l2 인 경우, 구간이 늘어났으므로 노드 삭제
}
 
void dfs(int now){
    visit[now] = true;
    st[now] = arr_n+1;
    arr[++arr_n] = now;
    for(auto next: adj[now]){
        if(!visit[next]){
            dfs(next);
        }
    }
    en[now] = arr_n+1;
    arr[++arr_n] = now;
}
 
void add_edge(int v1,int v2){
    adj[v1].push_back(v2);
    adj[v2].push_back(v1);
}
 
 
const int LOG_N = 22;
vector<PII> adj2[N];     // first : 목적지, second : 거리
PII parent[LOG_N+2][N]; // parent[i][j] : j 와 level차가 2^i인 조상
bool chk[N];
int level[N];
long long length[N]; // root와의 거리
 
void LCA_dfs(int now) {
    chk[now] = true;
    for (int i = 0; i < adj2[now].size(); i++) {
        int next = adj2[now][i].first;
        long long dist = adj2[now][i].second;
        if (!chk[next]) {
            parent[0][next] = PII(now, dist);
            int ancestor = now;
            for (int j = 0; ancestor != 0; j++) {
                parent[j + 1][next] = PII(parent[j][ancestor].first, parent[j][next].second+parent[j][ancestor].second);
                ancestor = parent[j][ancestor].first;
            }
            length[next] = length[now] + dist;
            level[next] = level[now] + 1;
            LCA_dfs(next);
        }
    }
}
 
int LCA(int v1, int v2) {
    if (level[v1] < level[v2])
        swap(v1, v2);
    for (int i = LOG_N; i >= 0; i--)
        if (((level[v1] - level[v2]) >> i) & 1)
            v1 = parent[i][v1].first;
    if (v1 == v2) return v1;
    for (int i = LOG_N; i >= 0; i--) {
        if (parent[i][v1].first != parent[i][v2].first) {
            v1 = parent[i][v1].first;
            v2 = parent[i][v2].first;
        }
    }
    return parent[0][v1].first;
}
 
void add_edge2(int v1,int v2,long long c){
    adj2[v1].push_back(PII(v2, c));
    adj2[v2].push_back(PII(v1, c));
 
}
 
void default_input(){
    scanf("%d",&n);
    for(int i=1; i<=n; i++) scanf("%d",&a[i]);
    for(int i=1; i<n; i++){
        int v1,v2;
        scanf("%d %d",&v1,&v2);
        add_edge(v1,v2);
        add_edge2(v1-1,v2-1,1);
    }
    dfs(1);
    LCA_dfs(0);
    scanf("%d",&q);
    for(int i=1; i<=q; i++){
        int x,y;
        scanf("%d %d",&x, &y);
        if(st[x]>st[y]) swap(x,y);
        if(LCA(x-1,y-1)==x-1)
            query[i] = PII(st[x],st[y]);
        else{
            lc[i] = 1+LCA(x-1,y-1);
            query[i] = PII(en[x],st[y]);
        }
        order[i].first = PII(query[i].first/SQRT_N, query[i].second);
        order[i].second = i;
    }
    sort(order+1, order+q+1);
    for(int i=1; i<=q; i++){
        go_query(query[order[i-1].second], query[order[i].second]);
        int lc2 = lc[order[i].second];
        if(lc2!=0) add2(lc2);
        ans[order[i].second] = an;
        if(lc2!=0) erase2(lc2);
    }
    for(int i=1; i<=q; i++)
        printf("%d\n",ans[i]);
}
 
int main(){
    default_input();
    return 0;
}
Colored by Color Scripter