HDU 4358 Boring counting 树状数组+思路

研究了整整一天orz……直接上官方题解神思路

#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <vector>
#include <algorithm>

using namespace std;

const int MAXN = 100100;

struct node
{
    int v, next;
};

struct subTree
{
    int st, ed;
};

struct Queryy
{
    int i;
    int st, ed;
};

int N, K, Q;
int EdgeN, TimeFlag;
node D[MAXN];   //树节点
int C[MAXN];    //树状数组
int head[MAXN];
subTree tree[MAXN];   //子树映射成线性序列之后对应的区间
Queryy qry[MAXN];     //查询信息
int weight[MAXN], addr[MAXN];
int ans[MAXN];
vector<int> pos[MAXN];

void AddEdge( int u, int v )
{
    D[EdgeN].v = v;
    D[EdgeN].next = head[u];
    head[u] = EdgeN++;
    return;
}

void DFS( int cur )                    //树形结构转线性结构
{
    tree[ cur ].st = ++TimeFlag;
    addr[ TimeFlag ] = weight[ cur ];
    for ( int i = head[cur]; i != -1; i = D[i].next )
        DFS( D[i].v );
    tree[cur].ed = TimeFlag;
    return;
}

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

void add( int x, int val )
{
    while ( x <= N )
    {
        C[x] += val;
        x += lowbit(x);
    }
    return;
}

int query( int x )
{
    int res = 0;
    while ( x > 0 )
    {
        res += C[x];
        x -= lowbit(x);
    }
    return res;
}

bool cmp( int a, int b )
{
    return weight[a] < weight[b];
}

void init()    //将weight离散化
{
    sort( addr + 1, addr + N + 1, cmp );

    int cnt = 0, pre = -1;
    for ( int i = 1; i <= N; ++i )
    {
        if ( weight[ addr[i] ] != pre )
            pre = weight[ addr[i] ], weight[ addr[i] ] = ++cnt;
        else weight[ addr[i] ] = cnt;
    }
    return;
}

bool cmp2( Queryy a, Queryy b )
{
    return a.ed < b.ed;
}

void solved()
{
    for ( int i = 0; i <= N; ++i ) pos[i].clear();

    sort( qry, qry + Q, cmp2 );
    int cur = 0;
    for ( int i = 1; i <= N; ++i )
    {
        int val = addr[i];
        pos[val].push_back(i);
        int sz = pos[val].size();
        if ( sz == K )
            add( pos[val][ sz - K ], 1 );
        else if ( sz > K )
        {
            add( pos[val][ sz - K ], 1 );
            add( pos[val][ sz - K - 1 ], -2 );
        }
        //printf( "ed = %d
", qry[cur].ed );
        while ( cur < Q && qry[cur].ed == i )
        {
            int id = qry[cur].i;
            ans[id] = query( qry[cur].ed ) - query( qry[cur].st - 1 );
          //  printf("ans[%d] = %d
", id, ans[id] );
            ++cur;
        }
    }
    return;
}

int main()
{
    int T, cas = 0;
    scanf( "%d", &T );
    while ( T-- )
    {
        memset( head, -1, sizeof( head ) );
        memset( C, 0, sizeof(C) );

        scanf( "%d%d", &N, &K );
        for ( int i = 1; i <= N; ++i )
        {
            scanf( "%d", &weight[i] );
            addr[i] = i;
        }
        init();

        EdgeN = 0;

        for ( int i = 1; i < N; ++i )      //建树
        {
            int u, v;
            scanf( "%d%d", &u, &v );
            AddEdge( u, v );
        }

        TimeFlag = 0;
        DFS(1);                           //树形结构转线性结构

        scanf( "%d", &Q );
        printf( "Case #%d:
", ++cas );
        for ( int i = 0; i < Q; ++i )
        {
            int u;
            scanf( "%d", &u );
            qry[i].i = i;
            qry[i].st = tree[u].st;
            qry[i].ed = tree[u].ed;
           // printf( "%d %d
", qry[i].st, qry[i].ed );
        }
        solved();
        for ( int i = 0; i < Q; ++i )
            printf( "%d
", ans[i] );

        if ( T ) puts("");
    }
    return 0;
}