CodeForces 375D Tree and Queries 莫队算法

题目：http://codeforces.com/problemset/problem/375/D

题意：给定一个以1为根的树，树中每个节点有一个颜色，求以某个节点v为根的子树中颜色出现次数大于等于k的颜色有几种

思路：dfs序分块，然后莫队乱搞，统计每种颜色的次数，用树状数组维护次数的前缀和，查询大于等于k次的颜色便是sum(n) - sum(k-1)

#include 
#include 
#include 
#include 
#include 
#define debug() puts("here")
using namespace std;

const int N = 100010;
struct edge
{
    int to, next;
}g[N*2];
struct node
{
    int v, l, r, k, id;
}q[N];
int n, m, unit, tmp;
int cnt, head[N];
int tot, a[N], in[N], out[N], num[N], val[N], res[N], bit[N];
void add_edge(int v, int u)
{
    g[cnt].to = u, g[cnt].next = head[v], head[v] = cnt++;
}
void dfs(int v, int fa)
{
    in[v] = ++tot;
    val[tot] = a[v];
    for(int i = head[v]; i != -1; i = g[i].next)
    {
        int u = g[i].to;
        if(u != fa) dfs(u, v);
    }
    out[v] = tot;
}
bool cmp(node a, node b)
{
    return a.l/unit != b.l/unit ? a.l/unit < b.l/unit : a.r < b.r;
}
void add(int i, int x)
{
    while(i <= n) bit[i] += x, i += i & -i;
}
int sum(int i)
{
    int res = 0;
    while(i > 0) res += bit[i], i -= i & -i;
    return res;
}
void _add(int i)
{
    if(num[i] > 0) add(num[i], -1);
    num[i]++;
    add(num[i], 1);
}
void del(int i)
{
    add(num[i], -1);
    num[i]--;
    if(num[i] > 0) add(num[i], 1);
}
void solve()
{
    dfs(1, -1);
    unit = (int)sqrt(1.0*n);
    for(int i = 1; i <= m; i++)
        q[i].id = i, q[i].l = in[q[i].v], q[i].r = out[q[i].v];
    sort(q+1, q+1+m, cmp);
    int l = 1, r = 0;
    tmp = 0;
    for(int i = 1; i <= m; i++)
    {
        while(r < q[i].r) _add(val[++r]);
        while(r > q[i].r) del(val[r--]);
        while(l < q[i].l) del(val[l++]);
        while(l > q[i].l) _add(val[--l]);
        if(q[i].k > n) res[q[i].id] = 0;
        else res[q[i].id] = sum(n) - sum(q[i].k - 1);
    }
    for(int i = 1; i <= m; i++) printf("%d\n", res[i]);
}
int main()
{
    scanf("%d%d", &n, &m);
    for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
    cnt = 0;
    memset(head, -1, sizeof head);
    int v, u;
    for(int i = 1; i <= n - 1; i++)
    {
        scanf("%d%d", &v, &u);
        add_edge(v, u), add_edge(u, v);
    }
    for(int i = 1; i <= m; i++) scanf("%d%d", &q[i].v, &q[i].k);
    solve();
    return 0;
}

另外一种不用树状数组的写法，直接用一个数组维护

#include 
#include 
#include 
#include 
#include 
#define debug() puts("here")
using namespace std;

const int N = 100010;
struct edge
{
    int to, next;
}g[N*2];
struct node
{
    int v, l, r, k, id;
}q[N];
int n, m, unit, tmp;
int cnt, head[N];
int tot, a[N], in[N], out[N], num[N], val[N], res[N], bit[N], rec[N];
void add_edge(int v, int u)
{
    g[cnt].to = u, g[cnt].next = head[v], head[v] = cnt++;
}
void dfs(int v, int fa)
{
    in[v] = ++tot;
    val[tot] = a[v];
    for(int i = head[v]; i != -1; i = g[i].next)
    {
        int u = g[i].to;
        if(u != fa) dfs(u, v);
    }
    out[v] = tot;
}
bool cmp(node a, node b)
{
    return a.l/unit != b.l/unit ? a.l/unit < b.l/unit : a.r < b.r;
}
void add(int i)
{
    rec[++num[i]]++;
}
void del(int i)
{
    rec[num[i]--]--;
}
void solve()
{
    tot = 0;
    dfs(1, -1);
    unit = (int)sqrt(1.0*n);
    for(int i = 1; i <= m; i++)
        q[i].id = i, q[i].l = in[q[i].v], q[i].r = out[q[i].v];
    sort(q+1, q+1+m, cmp);
    int l = 1, r = 0;
    tmp = 0;
    for(int i = 1; i <= m; i++)
    {
        while(r < q[i].r) add(val[++r]);
        while(r > q[i].r) del(val[r--]);
        while(l < q[i].l) del(val[l++]);
        while(l > q[i].l) add(val[--l]);
        res[q[i].id] = rec[q[i].k];
    }
    for(int i = 1; i <= m; i++) printf("%d\n", res[i]);
}
int main()
{
    scanf("%d%d", &n, &m);
    for(int i = 1; i <= n; i++) scanf("%d", &a[i]);
    cnt = 0;
    memset(head, -1, sizeof head);
    int v, u;
    for(int i = 1; i <= n - 1; i++)
    {
        scanf("%d%d", &v, &u);
        add_edge(v, u), add_edge(u, v);
    }
    for(int i = 1; i <= m; i++) scanf("%d%d", &q[i].v, &q[i].k);
    solve();
    return 0;
}