백준26618 정기 모임 4

작성일 | In

문제 링크

  • https://icpc.me/26618

사용 알고리즘

  • 센트로이드

시간복잡도

  • $O(N \log^2 N)$

풀이

간선을 반으로 쪼개서 정점을 하나씩 새로 만들어 줍시다. 어떤 정점과 거리가 $D_i$ 만큼 떨어져 있는 정점의 가중치 합을 구하는 문제입니다.

이런 문제는 보통 쿼리를 오프라인으로 받은 다음, 센트로이드 트리에서 문제를 해결할 수 있습니다. 정점 $x$가 $v$와 정확히 $d$ 만큼 떨어져 있기 위해서는, 센트로이드 트리 상에서 $v$의 조상 $p$가 $x$와 $d-dist(v, p)$ 만큼 떨어져 있으면서 원본 트리 상에서 $x - p$ 경로 위에 $v$가 없어야 합니다. 즉, 센트로이드 트리 상에서 $p$를 기준으로 봤을 때 $v$와 $x$가 서로 다른 서브 트리에 존재해야 합니다.
이런 식으로 서로 다른 서브 트리에서 오는 정점을 처리하는 것은 BOJ 13513 트리와 쿼리 4를 통해 연습할 수 있습니다.

센트로이드 트리의 각 정점에서 자신과 $d$ 만큼 떨어진 정점들의 가중치의 합을 저장하는 std::map 등을 관리하면 $O(N \log^2 N)$ 정도 시간에 문제를 해결할 수 있습니다.

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
#include <bits/stdc++.h>
using namespace std;

int N, Q, S[202020], U[202020], CPar[202020], CPos[202020];
vector<pair<int,int>> G[202020];
vector<int> CChild[202020];
vector<map<int,int>> CTree[202020];
map<int,int> Tree[202020];
int P[22][202020], D[202020], C[202020];

void DFS(int v, int b=-1){
    for(auto [i,w] : G[v]) if(i != b) P[0][i] = v, D[i] = D[v] + 1, C[i] = C[v] + w, DFS(i, v);
}

int LCA(int u, int v){
    if(D[u] < D[v]) swap(u, v);
    int diff = D[u] - D[v];
    for(int i=0; diff; i++, diff>>=1) if(diff & 1) u = P[i][u];
    if(u == v) return u;
    for(int i=21; i>=0; i--) if(P[i][u] != P[i][v]) u = P[i][u], v = P[i][v];
    return P[0][u];
}

int Dist(int u, int v){
    return C[u] + C[v] - 2 * C[LCA(u, v)];
}

int GetSize(int v, int b=-1){
    S[v] = 1;
    for(auto [i,w] : G[v]) if(i != b && !U[i]) S[v] += GetSize(i, v);
    return S[v];
}

int GetCent(int v, int sz, int b=-1){
    for(auto [i,w] : G[v]) if(i != b && !U[i] && S[i] * 2 > sz) return GetCent(i, sz, v);
    return v;
}

void Gather(int v, int b, int d, vector<int> &list){
    if(v <= N) list.push_back(d);
    for(auto [i,w] : G[v]) if(i != b && !U[i]) Gather(i, v, d+w, list);
}

int Build(int v){
    v = GetCent(v, GetSize(v)); U[v] = 1;
    int ch = count_if(G[v].begin(), G[v].end(), [](auto t){ return U[t.first] == 0; }) + 1;
    CTree[v].resize(ch); CChild[v].resize(ch); ch = 1;
    if(v <= N) Tree[v][0]++, CTree[v][0][0]++;
    for(auto [i,w] : G[v]){
        if(U[i]) continue;
        vector<int> vec; Gather(i, v, w, vec);
        for(auto d : vec) Tree[v][d]++, CTree[v][ch][d]++;
        ch++;
    }
    CChild[v][0] = v; ch = 1;
    for(auto [i,w] : G[v]){
        if(U[i]) continue;
        int now = Build(i);
        CChild[v][ch] = now;
        CPar[now] = v; CPos[now] = ch++;
    }
    return v;
}

int Query(int v, int d){
    int res = 0;
    for(int i=v, j=-1; i; j=CPos[i], i=CPar[i]){
        // s + x = d, x = d-s
        int s = Dist(v, i);
        if(auto it=Tree[i].find(d-s); it != Tree[i].end()) res += it->second;
        if(j != -1){
            auto it = CTree[i][j].find(d-s);
            if(it != CTree[i][j].end()) res -= it->second;
        }
        j = CPos[i];
    }
    return res;
}

int main(){
    ios_base::sync_with_stdio(false); cin.tie(nullptr);
    cin >> N >> Q;
    for(int i=1; i<N; i++){
        int a, b, c; cin >> a >> b >> c;
        G[a].emplace_back(N+i, c/2); G[N+i].emplace_back(a, c/2);
        G[b].emplace_back(N+i, c/2); G[N+i].emplace_back(b, c/2);
    }
    int Root = Build(1);
    DFS(1);
    for(int i=1; i<22; i++) for(int j=1; j<N+N; j++) P[i][j] = P[i-1][P[i-1][j]];
    for(int q=1; q<=Q; q++){
        int v, d; cin >> v >> d;
        cout << Query(N+v, d) << "\n";
    }
}