TcgE95 - Online C++0x Compiler & Debugging Tool

/*
* @Author: hungeazy
* @Date:   2025-11-26 13:25:55
* @Last Modified by:   hungeazy
* @Last Modified time: 2025-11-26 15:07:53
*/
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
// #pragma GCC optimize("O3")
// #pragma GCC optimize("unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,popcnt,lzcnt")
using namespace std;
using namespace __gnu_pbds;
bool M1;
#define fast ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
#define int long long
#define ll long long
#define ull unsigned long long
#define sz(x) x.size()
#define sqr(x) (1LL * (x) * (x))
#define all(x) x.begin(), x.end()
#define fill(f,x) memset(f,x,sizeof(f))
#define FOR(i,l,r) for(int i=l;i<=r;i++)
#define FOD(i,r,l) for(int i=r;i>=l;i--)
#define debug(x) cout << #x << " = " << x << '\n'
#define ii pair<int,int>
#define iii pair<int,ii>
#define di pair<ii,ii>
#define vi vector<int>
#define vii vector<ii>
#define mii map<int,int>
#define fi first
#define se second
#define pb push_back
#define MOD 1000000007
#define __lcm(a,b) (1ll * ((a) / __gcd((a), (b))) * (b))
#define YES cout << "YES\n"
#define NO cout << "NO\n"
#define MASK(i) (1LL << (i))
#define c_bit(i) __builtin_popcountll(i)
#define BIT(x,i) ((x) & MASK(i))
#define SET_ON(x,i) ((x) | MASK(i))
#define SET_OFF(x,i) ((x) & ~MASK(i))
#define oo 1e18
#define name "DRAGONFLY"
#define endl '\n'
#define memory() cerr << abs(&M2-&M1)/1024.0/1024 << " MB" << endl
#define time() cerr << endl << "-------------Time:" << 1000.0 * clock() / CLOCKS_PER_SEC << "ms." << endl
template<typename T> bool maximize(T &res, const T &val) { if (res < val){ res = val; return true; }; return false; }
template<typename T> bool minimize(T &res, const T &val) { if (res > val){ res = val; return true; }; return false; }
template <class T> using ordered_set = tree <T, null_type, less_equal <T>, rb_tree_tag,tree_order_statistics_node_update>;
const int N = (int)2e5+10;
int n,d,b[N],s[N],pos[N];
vi g[N];
ii edge[N];

namespace sub1 {

    int par[N];
    bool check[N];

    bool approved() {
        return n <= 1e3 and d <= 1e3;
    }

    void DFS(int u, int p)
    {
        for (int v : g[u])
            if (v != p)
            {
                par[v] = u;
                DFS(v,u);
            }
    }

    void solve(void)
    {
        DFS(1,-1);
        FOR(i,1,d)
        {
            vi vec;
            int u = pos[i];
            while (u != 1)
            {
                vec.pb(u);
                u = par[u];
            }
            vec.pb(1);
            FOR(j,1,n) check[j] = false;
            reverse(all(vec));
            int ans = 0;
            for (int x : vec)
                if (b[x] > 0)
                {
                    b[x]--;
                    if (!check[s[x]])
                    {
                        check[s[x]] = true;
                        ans++;
                    }
                }
            cout << ans << " ";
        }
    }

}

namespace sub2 {

    int lastSeen[N],ans[N],par[N];

    bool approved() {
        FOR(i,1,n)
            if (b[i] != d) return false;
        return true;
    }

    void DFS(int u, int p)
    {
        for (int v : g[u])
            if (v != p)
            {
                par[v] = u;
                DFS(v,u);
            }
    }

    void solve(void)
    {
        DFS(1,-1);
        FOR(i,1,n)
        {
            int u = i, dist = 0;
            while (u != 1)
            {
                if (lastSeen[s[u]] != i)
                {
                    lastSeen[s[u]] = i;
                    dist++;
                }
                u = par[u];
            }
            if (lastSeen[s[1]] != i)
            {
                lastSeen[s[1]] = i;
                dist++;
            }
            ans[i] = dist;
        }
        FOR(i,1,d) cout << ans[pos[i]] << " ";
    }

}

namespace sub6 {

    int par[N],h[N],nxt[N],lastSeen[N],ans[N];

    void DFS(int u, int p)
    {
        for (int v : g[u])
            if (v != p)
            {
                par[v] = u;
                h[v] = h[u]+1;
                DFS(v,u);
            }
    }

    int findNxt(int u)
    {
        if (!u) return 0;
        if (nxt[u] == u) {
            if (b[u] > 0) return u;
        }
        if (nxt[u] == u) nxt[u] = par[u];
        return (nxt[u] == u ? u : nxt[u] = findNxt(nxt[u]));
    }

	void solve(void)
	{
        DFS(1,-1);
        FOR(i,1,n) nxt[i] = i;
        FOR(i,1,d)
        {
            int distinct = 0, cur = pos[i];
            while (cur)
            {
                int u = findNxt(cur);
                // debug(cur); debug(u);
                if (!u) break;
                if (lastSeen[s[u]] != i)
                {
                    lastSeen[s[u]] = i;
                    distinct++;
                }
                b[u]--;
                if (!b[u]) nxt[u] = par[u];
                cur = par[u];
            }
            ans[i] = distinct;
        }
        FOR(i,1,d) cout << ans[i] << " ";
	}

}

bool M2;
signed main()
{
    fast;
    if (fopen(name".inp","r"))
    {
    	freopen(name".inp","r",stdin);
    	freopen(name".out","w",stdout);
    }
    cin >> n >> d;
    FOR(i,1,n) cin >> b[i];
    FOR(i,1,n) cin >> s[i];
    FOR(i,1,d) cin >> pos[i];
    FOR(i,1,n-1)
    {
    	int u,v;
    	cin >> u >> v;
    	g[u].pb(v); g[v].pb(u);
        edge[i] = {u,v};
    }
    if (sub1::approved()) return sub1::solve(), time(), memory(), 0;
    if (sub2::approved()) return sub2::solve(), time(), memory(), 0;
    sub6::solve();
    time();
    memory();
    return 0;
}
// ██░ ██  █    ██  ███▄    █   ▄████
//▓██░ ██▒ ██  ▓██▒ ██ ▀█   █  ██▒ ▀█▒
//▒██▀▀██░▓██  ▒██░▓██  ▀█ ██▒▒██░▄▄▄░
//░▓█ ░██ ▓▓█  ░██░▓██▒  ▐▌██▒░▓█  ██▓
//░▓█▒░██▓▒▒█████▓ ▒██░   ▓██░░▒▓███▀▒
// ▒ ░░▒░▒░▒▓▒ ▒ ▒ ░ ▒░   ▒ ▒  ░▒   ▒
// ▒ ░▒░ ░░░▒░ ░ ░ ░ ░░   ░ ▒░  ░   ░
// ░  ░░ ░ ░░░ ░ ░    ░   ░ ░ ░ ░   ░
// ░  ░  ░   ░              ░       ░