Algoogle

Algorithm for Programming Contest

Codeforces 375D Tree and Queries

Category: Codeforces Tag: data-structure, sqrt-decomposition, tree

Tree and Queries

問題概要


N(<=100000)頂点の木の各頂点に色が塗られている. このとき以下のクエリに答えろ.
頂点vの部分木にk個以上に塗られている色はいくつあるか

解法


オイラーツアーして平方分割する.
部分木の開始点を平方分割の各バケットの要素数Bごとに区切る.
クエリをその開始点順にソートし, 同じバケットでは終了点順でソートしておくと開始点の各バケットでの移動は高々B回(オイラーツアーをしているので).
また各バケットでのクエリの終了点の移動は高々N回になる.
よって全体ではになる.

データ構造をマージする一般的なテクも使えるらしいのでコードを載せておいた.
深い順にクエリを処理することで順にマージしていけるようにする.
マージは大きいものを選んでそれに小さい方を愚直に突っ込んでいるだけ(直感的にはすごくTLEしそうな感じだ).
マージ後の小さい方はclearしておかないとたぶんMLEする.

コード


(375D.cpp) download
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
#include <bits/stdc++.h>
using namespace std;
#define repi(i,a,b) for(int i = (int)(a); i < (int)(b); i++)
#define rep(i,a) repi(i,0,a)
#define all(u) begin(u),end(u)
#define pb push_back
const int M = 100010, B = 400;
int n, m, idx, par[M],
    ans[M], l[M], r[M], id[M], // [l, r)
    c[M], cc[M], ccc[M]; // color, color count, color count count
struct query{
    int v, k, id;
    bool operator<(const query &q)const{
        if(l[v]/B != l[q.v]/B) return l[v] < l[q.v];
        return r[v] < r[q.v];
    }
};

vector<int> G[M];
query qs[M];

void build_lr(int v, int u){
    id[idx] = v;
    par[v] = u;
    l[v] = idx++;
    for(int &w: G[v])
        if(w != u) build_lr(w,v);
    r[v] = idx;
}


void solve(){
    build_lr(0,0);
    sort(qs, qs+m);
    int ls = 0, rs = 0;
    rep(i,m){
        query &q = qs[i];
        while(ls < l[q.v]) --ccc[cc[c[id[ls++]]]--];
        while(ls > l[q.v]) ++ccc[++cc[c[id[--ls]]]];
        while(rs < r[q.v]) ++ccc[++cc[c[id[rs++]]]];
        while(rs > r[q.v]) --ccc[cc[c[id[--rs]]]--];
        ans[q.id] = ccc[q.k];
    }
    rep(i,m) printf("%d\n", ans[i]);
}

void input(){
    scanf("%d%d", &n, &m);
    rep(i,n) scanf("%d", c+i);
    rep(i,n-1){
        int u, v;
        scanf("%d%d", &u, &v);
        u--; v--;
        G[v].pb(u);
        G[u].pb(v);
    }
    rep(i,m){
        int v, k;
        scanf("%d%d", &v, &k);
        qs[i] = (query){v-1,k,i};
    }
}

int main(){
    input();
    solve();
    return 0;
}
(375D1.cpp) download
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
96
97
98
99
100
101
102
103
104
105
106
#include <bits/stdc++.h>
using namespace std;
#define repi(i,a,b) for(int i = (int)(a); i < (int)(b); i++)
#define rep(i,a) repi(i,0,a)
#define all(u) begin(u),end(u)
#define pb push_back

const int M = 100010;
int n, m, idx, depth[M], par[M], ans[M], c[M];

struct query{
    int v, k, id;
    bool operator<(const query &q)const{
        return depth[v] > depth[q.v];
    }
};

vector<int> G[M], ccc[M];
query qs[M];
map<int,int> cnt[M];

void build(int v, int u){
    par[v] = u;
    depth[v] = depth[u]+1;
    for(int &w: G[v])
        if(w != u) build(w,v);
}

// size of v >= size of u
void merge(int v, int u){
    for(auto &p: cnt[u]){
        int &cvf = cnt[v][p.first];
        repi(i,cvf+1,cvf+p.second+1){
            if((int)ccc[v].size() == i) ccc[v].pb(1);
            else ccc[v][i]++;
        }
        cvf += p.second;
    }
    cnt[u].clear();
    ccc[u].clear();
}

void dfs(int v){
    if(cnt[v].size()) return;
    int id = -1, w = 0;
    for(int &u: G[v])
        if(u != par[v]){
            dfs(u);
            if((int)cnt[u].size() > w){
                w = cnt[u].size();
                id = u;
            }
        }
    if(id < 0){
        cnt[v][c[v]] = 1;
        ccc[v] = {0,1};
        return;
    }
    swap(cnt[v], cnt[id]);
    swap(ccc[v], ccc[id]);
    for(int &u: G[v]){
        if(u != par[v] and u != id)
            merge(v, u);
    }
    int &cv = cnt[v][c[v]];
    if((int)ccc[v].size() == cv+1) ccc[v].pb(1);
    else ccc[v][cv+1]++;
    cv++;
}

void solve(){
    depth[0] = -1;
    build(0,0);
    sort(qs, qs+m);
    rep(i,m){
        query &q = qs[i];
        dfs(q.v);
        if(q.k >= (int)ccc[q.v].size()) continue;
        ans[q.id] = ccc[q.v][q.k];
    }
    rep(i,m) printf("%d\n", ans[i]);
}

void input(){
    scanf("%d%d", &n, &m);
    rep(i,n) scanf("%d", c+i);
    rep(i,n-1){
        int u, v;
        scanf("%d%d", &u, &v);
        u--; v--;
        G[v].pb(u);
        G[u].pb(v);
    }
    rep(i,m){
        int v, k;
        scanf("%d%d", &v, &k);
        qs[i] = (query){v-1,k,i};
    }
}


int main(){
    input();
    solve();
    return 0;
}

Comments