zstu4143 Tree Query (函数式线段树)

4143: Tree Query

代码：

#include <stdio.h>
#include <string.h>
#include <algorithm>
#include <map>
#include <math.h>
#include <queue>
#include <vector>
#include <string>
#include <iostream>
#include <stdlib.h>
#include <time.h>
#define lowbit(x) (x&(-x))
#define ll __int64
#define lson l , m , rt << 1
#define rson m + 1 , r , rt << 1 | 1
#define ls son[0][rt]
#define rs son[1][rt]
#define new_edge(a,b,c) edge[tot].t = b , edge[tot].v = c , edge[tot].next = head[a] , head[a] = tot ++
using namespace std;

const int maxn = 111111 ;
int L[maxn<<5] , R[maxn<<5] , cnt[maxn<<5] , sum[maxn<<5] ;
int root[maxn] , tail , tim ;
int num[maxn] , a[maxn] ;
struct Point {
    int l , r ;
} p[maxn] ;

struct Edge {
    int t , next , v ;
} edge[maxn<<2] ;
int head[maxn] , tot ;
int N = maxn ;

void push_up ( int rt ) {
    sum[rt] = sum[L[rt]] + sum[R[rt]] ;
    cnt[rt] = cnt[L[rt]] + cnt[R[rt]] ;
}

int build ( int l , int r ) {
    int rt = ++ tail ;
    sum[rt] = cnt[rt] = 0 ;
    if ( l == r ) return rt ;
    int m = l + r >> 1 ;
    L[rt] = build ( l , m ) ;
    R[rt] = build ( m + 1 , r ) ;
    return rt ;
}

void init () {
    memset ( head , -1 , sizeof ( head ) ) ;
    tim = tail = tot = 0 ;
    root[0] = build ( 1 , N ) ;
}

void dfs ( int u , int fa ) {
    tim ++ ;
    p[u].l = tim ; a[tim] = num[u] ;
    int i ;
    for ( i = head[u] ; i != -1 ; i = edge[i].next ) {
        int e = edge[i].t ;
        if ( e == fa ) continue ;
        dfs ( e , u ) ;
    }
    p[u].r = tim ;
}

int update ( int fa , int p , int l , int r ) {
    int rt = ++ tail ;
    if ( l == r ) {
        sum[rt] = sum[fa] + p ;
        cnt[rt] = cnt[fa] + 1 ;
        return rt ;
    }
    int m = l + r >> 1 ;
    L[rt] = L[fa] , R[rt] = R[fa] ;
    if ( p <= m ) L[rt] = update ( L[fa] , p , l , m ) ;
    else R[rt] = update ( R[fa] , p , m + 1 , r ) ;
    push_up ( rt ) ;
    return rt ;
}

int query ( int u , int v , int l , int r , int c , int w ) {
    if ( l == r ) {
        int add = cnt[L[v]] - cnt[L[u]] ;
        return c + min ( w / l , add ) ;
    }
    int m = l + r >> 1 ;
    int num = sum[L[v]] - sum[L[u]] ;
    int add = cnt[L[v]] - cnt[L[u]] ;
    if ( num >= w ) return query ( L[u] , L[v] , l , m , c , w ) ;
    else return query ( R[u] , R[v] , m + 1 , r , c + add , w - num ) ;
}

int main() {
    int n , m , i , j , k ;
    while ( scanf ( "%d%d" , &n , &m ) == 2 ) {
        for ( i = 1 ; i <= n ; i ++ ) scanf ( "%d" , &num[i] ) ;
        init () ;
        for ( i = 1 ; i < n ; i ++ ) {
            int a , b ;
            scanf ( "%d%d" , &a , &b ) ;
            new_edge ( a , b , 0 ) ;
            new_edge ( b , a , 0 ) ;
        }
        dfs ( 1 , 1 ) ;
        for ( i = 1 ; i <= n ; i ++ ) 
            root[i] = update ( root[i-1] , a[i] , 1 , N ) ;
        for ( i = 1 ; i <= n ; i ++ )
            printf ( "%d%c" , query ( root[p[i].l-1] , root[p[i].r] , 1 , N , 0 , m ) , i == n ? '\n' : ' ' ) ;
    }
    return 0;
}