树状数组+离散化+DFS序+离线 HDOJ 4358 Boring counting

题目传送门

题意：给你一棵树，树上的每个节点都有树值，给m个查询，问以每个点u为根的子树下有多少种权值恰好出现k次。

分析：首先要对权值离散化，然后要将树形转换为线形，配上图：

收获：

//还没写完。。。

代码：

/************************************************
* Author        :Running_Time
* Created Time  :2015/9/10 星期四 19:15:22
* File Name     :I.cpp
 ************************************************/

#include <cstdio>
#include <algorithm>
#include <iostream>
#include <sstream>
#include <cstring>
#include <cmath>
#include <string>
#include <vector>
#include <queue>
#include <deque>
#include <stack>
#include <list>
#include <map>
#include <set>
#include <bitset>
#include <cstdlib>
#include <ctime>
using namespace std;

#define lson l, mid, rt << 1
#define rson mid + 1, r, rt << 1 | 1
typedef long long ll;
const int N = 1e5 + 10;
const int INF = 0x3f3f3f3f;
const int MOD = 1e9 + 7;
struct Edge {
    int v, nex;
}edge[N*2];
struct Query    {
    int l, r, id;
    Query (int _l = 0, int _r = 0, int _id = 0) : l (_l), r (_r), id (_id) {}
    bool operator < (const Query &x) const  {
        return r < x.r;
    }
}q[N];
struct BIT  {
    int c[N];
    void init(void) {
        memset (c, 0, sizeof (c));
    }
    void updata(int i, int x)   {
        while (i < N)   {
            c[i] += x;  i += i & (-i);
        }
    }
    int query(int i)    {
        int ret = 0;
        while (i)   {
            ret += c[i];    i -= i & (-i);
        }
        return ret;
    }
}bit;
int head[N], dfn[N], low[N], w[N], p[N], val[N], ans[N];
vector<int> cnt[N];
int e, dep;

void init(void) {
    memset (head, -1, sizeof (head));
    e = 0;  dep = 0;
}

void add_edge(int u, int v) {
    edge[e].v = v;  edge[e].nex = head[u];
    head[u] = e++;
}

bool cmp(int i, int j)  {
    return w[i] < w[j];
}

void compress(int n)    {
    for (int i=1; i<=n; ++i)    p[i] = i;
    sort (p+1, p+1+n, cmp);
    int rk = 0, pre = w[p[1]] - 1;
    for (int i=1; i<=n; ++i)    {
        if (pre != w[p[i]]) {
            pre = w[p[i]];
            w[p[i]] = ++rk;
        }
        else    {
            w[p[i]] = rk;
        }
    }
}

void DFS(int u, int fa) {
    dfn[u] = ++dep; val[dep] = w[u];
    for (int i=head[u]; ~i; i=edge[i].nex)  {
        int v = edge[i].v;
        if (v == fa)    continue;
        DFS (v, u);
    }
    low[u] = dep;
}

int main(void)    {
    int T, cas = 0;  scanf ("%d", &T);
    while (T--) {
        if (cas)    puts ("");
        printf ("Case #%d:\n", ++cas);
        int n, k;   scanf ("%d%d", &n, &k);
        init ();

        for (int i=1; i<=n; ++i)    {
            scanf ("%d", &w[i]);
        }
        compress (n);                                           //离散化，升序排序，相同的还是相同的

        for (int u, v, i=1; i<n; ++i)   {
            scanf ("%d%d", &u, &v);
            add_edge (u, v);    add_edge (v, u);
        }
        DFS (1, -1);                                            //先序遍历，得到DFS序，树形变线形

        int m;  scanf ("%d", &m);
        for (int u, i=1; i<=m; ++i)    {
            scanf ("%d", &u);
            q[i] = Query (dfn[u], low[u], i);
        }
        sort (q+1, q+1+m);                                      //按照DFS序排序

        for (int i=1; i<=n; ++i)    cnt[i].clear ();
        int qu = 1; bit.init ();
        for (int i=1; i<=n; ++i)    {                           //按照dep深度从小到大
            int v = val[i];
            cnt[v].push_back (i);                               //表示离散后的相同的权值的个数
            int sz = cnt[v].size ();
            if (sz >= k)    {
                if (sz == k)    bit.updata (cnt[v][sz-k], 1);
                else    {
                    bit.updata (cnt[v][sz-k-1], -2);            //?
                    bit.updata (cnt[v][sz-k], 1);
                }
            }
            while (qu <= m && q[qu].r == i) {
                ans[q[qu].id] = bit.query (q[qu].r) - bit.query (q[qu].l - 1);
                qu++;
            }
        }

        for (int i=1; i<=m; ++i)    {
            printf ("%d\n", ans[i]);
        }
    }

    return 0;
}