Codeforces Round #221 (Div. 1) D Tree and Queries

Tree and Queries

询问 \(V_j\) 的子树中，有多少种颜色出现了 \(K_j\) 次

启发式合并

最直接的，树上启发式合并的同时维护颜色出现的次数，然后再拿一个数组记录一下出现了 \(i\) 次的颜色数量，储存在 \(sum_i\)

对于这个 \(sum\) 的维护，一开始觉得直接套个树状数组上去，复杂度 \(O(nlog^2n)\)，直接就莽过去了

后来发现可以直接维护，不需要树状数组，因为我们记录颜色出现的次数必然是从 \(0\) 然后一个个加到最大值 \(cnt_i\) 的过程，因此，这个遍历过程本身就不会缺少，直接将 \(sum_j\) 定义为出现了 \(j\) 次及以上的颜色数量，就会发现非常好维护

时间复杂度为 \(O(nlogn)\)

直接维护：\(O(nlogn)\)

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <vector>
using namespace std;
#define pii pair<int, int>
typedef long long ll;
const int maxn = 1e5 + 10;
vector<int>gra[maxn];
int col[maxn], cnt[maxn], fa[maxn];
int siz[maxn], hson[maxn], ans[maxn];
int L[maxn], R[maxn], tp = 0, rnk[maxn];
int sum[maxn];
vector<pii>qs[maxn];

void dfs1(int now, int pre)
{
    L[now] = ++tp;
    fa[now] = pre;
    siz[now] = 1;
    hson[now] = -1;
    rnk[tp] = now;
    for(int nex : gra[now])
    {
        if(nex == pre) continue;
        dfs1(nex, now);
        siz[now] += siz[nex];
        if(hson[now] == -1 || siz[nex] > siz[hson[now]])
            hson[now] = nex;
    }
    R[now] = tp;
}

inline void add(int now)
{
    cnt[col[now]]++;
    sum[cnt[col[now]]]++;
}

inline void del(int now)
{
    sum[cnt[col[now]]]--;
    cnt[col[now]]--;
}

void dfs2(int now, bool keep)
{
    for(int nex : gra[now])
        if(nex != fa[now] && nex != hson[now]) dfs2(nex, false);
    if(hson[now] != -1) dfs2(hson[now], true);
    for(int nex : gra[now])
    {
        if(nex == fa[now] || nex == hson[now]) continue;
        for(int i=L[nex]; i<=R[nex]; i++)
            add(rnk[i]);
    }
    add(now);
    for(auto [k, id] : qs[now])
        ans[id] = sum[k];
    if(!keep)
    {
        for(int i=L[now]; i<=R[now]; i++)
            del(rnk[i]);
    }
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    int n, m;
    cin >> n >> m;
    for(int i=1; i<=n; i++) cin >> col[i];
    for(int i=1; i<n; i++)
    {
        int a, b;
        cin >> a >> b;
        gra[a].push_back(b);
        gra[b].push_back(a);
    }
    for(int i=0; i<m; i++)
    {
        int a, b;
        cin >> a >> b;
        qs[a].push_back({b, i});
    }
    dfs1(1, 1);
    dfs2(1, true);
    for(int i=0; i<m; i++)
        cout << ans[i] << "\n";
    return 0;
}

树状数组维护：\(O(nlogn^2n)\)

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <vector>
using namespace std;
#define pii pair<int, int>
typedef long long ll;
const int maxn = 1e5 + 10;
vector<int>gra[maxn];
int col[maxn], cnt[maxn], fa[maxn];
int siz[maxn], hson[maxn], ans[maxn];
int L[maxn], R[maxn], tp = 0, rnk[maxn];
vector<pii>qs[maxn];
ll tr[maxn], tot = 0;
int n, m;

inline int lowbit(int x)
{
    return x & (-x);
}

void dfs1(int now, int pre)
{
    L[now] = ++tp;
    fa[now] = pre;
    siz[now] = 1;
    hson[now] = -1;
    rnk[tp] = now;
    for(int nex : gra[now])
    {
        if(nex == pre) continue;
        dfs1(nex, now);
        siz[now] += siz[nex];
        if(hson[now] == -1 || siz[nex] > siz[hson[now]])
            hson[now] = nex;
    }
    R[now] = tp;
}

void update(int now, ll val)
{
    for(int i=now; i<maxn; i+=lowbit(i))
        tr[i] += val;
}

ll query(int now)
{
    ll ans = 0;
    for(int i=now; i; i-=lowbit(i))
        ans += tr[i];
    return ans;
}

inline void add(int now)
{
    if(cnt[col[now]]) update(cnt[col[now]], -1);
    else tot++;
    cnt[col[now]]++;
    update(cnt[col[now]], 1);
}

inline void del(int now)
{
    update(cnt[col[now]]--, -1);
    if(cnt[col[now]]) update(cnt[col[now]], 1);
    else tot--;
}

void dfs2(int now, bool keep)
{
    for(int nex : gra[now])
        if(nex != fa[now] && nex != hson[now]) dfs2(nex, false);
    if(hson[now] != -1) dfs2(hson[now], true);
    for(int nex : gra[now])
    {
        if(nex == fa[now] || nex == hson[now]) continue;
        for(int i=L[nex]; i<=R[nex]; i++)
            add(rnk[i]);
    }
    add(now);
    for(auto [k, id] : qs[now])
        ans[id] = tot - query(k - 1);
    if(!keep)
    {
        for(int i=L[now]; i<=R[now]; i++)
            del(rnk[i]);
    }
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    cin >> n >> m;
    for(int i=1; i<=n; i++) cin >> col[i];
    for(int i=1; i<n; i++)
    {
        int a, b;
        cin >> a >> b;
        gra[a].push_back(b);
        gra[b].push_back(a);
    }
    for(int i=0; i<m; i++)
    {
        int a, b;
        cin >> a >> b;
        qs[a].push_back({b, i});
    }
    dfs1(1, 1);
    dfs2(1, true);
    for(int i=0; i<m; i++)
        cout << ans[i] << "\n";
    return 0;
}