UJGOI 2022 Day 3C

#include <bits/stdc++.h>
using namespace std;

//#define int long long
#define mp make_pair
#define pb push_back

#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("avx2")

//#pragma GCC optimize("O3,unroll-loops")
//#pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")

constexpr int INF = (int)1e9;
constexpr int N = (1 << 18);

vector<int> g[N];

struct HLD{
    int t[4*N];
    int n;
    void init(int sz){
        n = sz-1;
        return;
    }
    struct SegTreeMn{
        struct node{
            int mn,pos;
            int w;
            node(){}
            node(int x,int pos):pos(pos){
                mn = x;
                w = 0;
            }
            node(node l,node r){
                mn = min(l.mn,r.mn);
                if(mn == l.mn)
                    pos = l.pos;
                else
                    pos = r.pos;
                w = 0;
            }
        } t[4*N];
        int n;
        SegTreeMn(){}
        void init(int sz){
            n = sz-1;
            return;
        }
        void push(int v){
            t[v<<1].mn += t[v].w;
            t[v<<1|1].mn += t[v].w;
            t[v<<1].w += t[v].w;
            t[v<<1|1].w += t[v].w;
            t[v].w = 0;
            return;
        }
        void build(int v,int l,int r,vector<int>& order,vector<int>& a){
            if(l == r){
                t[v] = node(a[order[l]],l);
                return;
            }
            int m = (l+r)>>1;
            build(v<<1,l,m,order,a);
            build(v<<1|1,m+1,r,order,a);
            t[v] = node(t[v<<1],t[v<<1|1]);
            return;
        }
        void build(vector<int>& order,vector<int>& a){
            n = order.size()-1;
            build(1,0,n,order,a);
            return;
        }
        void upd(int v,int l,int r,int pos,int nw){
            if(l == r){
                t[v] = node(nw,pos);
                return;
            }
            push(v);
            int m = (l+r)>>1;
            if(pos <= m)
                upd(v<<1,l,m,pos,nw);
            else
                upd(v<<1|1,m+1,r,pos,nw);
            t[v] = node(t[v<<1],t[v<<1|1]);
            return;
        }
        void updSegment(int v,int l,int r,int tl,int tr){
            if(l > r || tl > tr)
                return;
            if(l == tl && r == tr){
                t[v].w--;
                t[v].mn--;
                return;
            }
            push(v);
            int m = (l+r)>>1;
            updSegment(v<<1,l,m,tl,min(tr,m));
            updSegment(v<<1|1,m+1,r,max(tl,m+1),tr);
            t[v] = node(t[v<<1],t[v<<1|1]);
            return;
        }
        void updSegment(int l,int r){
            updSegment(1,0,n,l,r);
            return;
        }
        void upd(int pos,int nw){
            upd(1,0,n,pos,nw);
            return;
        }
    }
    tMin;

    struct SegTree{
        int n;
        int tree[2*N];
        SegTree(){}
        void init(int sz){
            return;
        }
        void upd(int pos, int newval) { // arr[pos] := newval
            pos += N;
            tree[pos] = newval;
            pos >>= 1;
            while (pos > 0) {
                tree[pos] = tree[pos << 1] + tree[(pos << 1) | 1];
                pos >>= 1;
            }
        }
        int get(int l, int r) { // [l, r)
            r++;
            l += N;
            r += N;
            int ans = 0;
            while (l < r) {
                if (l & 1) {
                    ans += tree[l++];
                }
                if (r & 1) {
                    ans += tree[--r];
                }
                l >>= 1;
                r >>= 1;
            }
            return ans;
        }
    }
    tSum;

    int d[N],sz[N];
    int go[N];
    int parent[N];
    int tin[N],tout[N];
    int head[N];
    int p[N];

    void dfs1(int v,int pr = -1){ /// determine heavy paths
        sz[v] = 1;
        int mx = 0;
        go[v] = -1;
        for(auto& to : g[v]){
            if(to == pr)
                continue;
            dfs1(to,v);
            if(sz[to] > mx){
                mx = sz[to];
                go[v] = to;
            }
            sz[v] += sz[to];
        }
        return;
    }

    vector<int> order;

    void dfs2(int v,int h,int pr = 0){
        p[v] = pr;
        tin[v] = order.size();
        head[v] = h;
        order.pb(v);
        if(go[v] == -1)
            return;
        dfs2(go[v],h,v);
        for(auto& to : g[v]){
            if(to == pr || to == go[v])
                continue;
            dfs2(to,to,v);
        }
        tout[v] = order.size()-1;
        return;
    }

    int cnt[N];
    int b[N],s[N];
    vector<int> was[N];
    vector<int> sons[N];
    int pp[N];
    bool deleted[N];

    int find_parent(int x){
        if(deleted[pp[x]])
            pp[x] = find_parent(pp[x]);
        return pp[x];
    }

    void dfs(int v,int pr = -1){
        cnt[s[v]]++;
        pp[v] = 0;
        if(!was[s[v]].empty()){
            pp[v] = was[s[v]].back();
            sons[was[s[v]].back()].pb(v);
        }
        was[s[v]].pb(v);
        if(cnt[s[v]] == 1){
            tSum.upd(tin[v],1);
//            cerr << "added in " << v << "\n";
        }
        tMin.upd(tin[v],b[v]);
        for(auto& to : g[v]){
            if(to == pr)
                continue;
            dfs(to,v);
        }
        cnt[s[v]]--;
        was[s[v]].pop_back();
        return;
    }

    int get(int v){
        int ans = 0;
//        cerr << "v = " << v << "\n";
//        cerr << "h = " << head[v] << "\n";
        for(int h = head[v];; h = head[v]){
            ans += tSum.get(tin[h],tin[v]);
            tMin.updSegment(tin[h],tin[v]);
            if(h == 1)
                break;
            v = p[h];
        }
        return ans;
    }

    void Normalize(){
        do{
            auto p = tMin.t[1];
            if(p.mn > 0)
                break;
            int pos = p.pos;
            int v = order[pos];
//            cerr << "going to delete " << v << " " << p.mn << "\n";
            deleted[v] = true;
            tMin.upd(pos,INF);
            if(tSum.get(tin[v],tin[v]) == 1){
                tSum.upd(tin[v],0);
                for(auto& x : sons[v]){
                    if(!deleted[x]){
                        tSum.upd(tin[x],1);
//                        cerr << "added " << x << "\n";
                    }
                }
            } else {
                find_parent(v);
                int pr = pp[v];
                if(pr != 0){
                    if(sons[v].size() > sons[pr].size()){
                        for(auto& x : sons[pr])
                            sons[v].pb(x);
                        swap(sons[v],sons[pr]);
                    } else {
                        for(auto& x : sons[v])
                            sons[pr].pb(x);
                    }
                }
            }
            deleted[v] = true;
//            cerr << "deleted " << v << "\n";
        }while(true);
    }

    void init(int sz,vector<int>& B,vector<int>& S){
        n = sz-1;
        tMin.init(sz);
        tSum.init(sz);
        dfs1(1);
        dfs2(1,1);
        for(int i = 0; i < B.size(); i++){
            b[i] = B[i];
            s[i] = S[i];
        }
        dfs(1);
        Normalize();
    }

    int answer(int v){
        int ans = get(v);
        Normalize();
        return ans;
    }

} Solver;

void solve(){
    int n,d;
//    cin >> n >> d;
    scanf("%d %d",&n,&d);

    vector<int> b(n+1,0);
    vector<int> s(n+1,0);

    for(int i = 1; i <= n; i++){
//        cin >> b[i];
        scanf("%d",&b[i]);
    }
    for(int i = 1; i <= n; i++){
//        cin >> s[i];
        scanf("%d",&s[i]);
    }

    vector<int> queries;
    for(int i = 0; i < d; i++){
        int v;
//        cin >> v;
        scanf("%d",&v);
        queries.pb(v);
    }

    for(int i = 1; i < n; i++){
        int a,b;
//        cin >> a >> b;
        scanf("%d %d",&a,&b);
        g[a].pb(b);
        g[b].pb(a);
    }

    Solver.init(n,b,s);
    for(auto& v : queries){
//        cerr << "Next QUERY!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n";
        int ans = Solver.answer(v);
//        cerr << "ANSWER = " << ans << "\n\n";
        printf("%d ",ans);
    }

    return;
}

signed main(){
    ios::sync_with_stdio(false);cin.tie(nullptr);cout.tie(nullptr);
    int tests = 1;
//    cin >> tests;
    for(int test = 1; test <= tests; test++){
//        cerr << "test = " << test << "\n";
        solve();
    }
    return 0;
}
/**
O(Nlog^2)

5 1
4 1 0 3 1
1 3 2 2 1
2
5 2
2 1
1 4
1 3

**/