Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include<bits/stdc++.h>
- #include<ext/pb_ds/assoc_container.hpp>
- #include<ext/pb_ds/tree_policy.hpp>
- using namespace __gnu_pbds;
- using namespace std;
- #define IO ios_base::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL);
- #define pb push_back
- #define all(x) (x).begin(),(x).end()
- #define ll long long
- #define ull unsigned long long
- #define ld long double
- #define eps 1e-9
- #define sz(a) (ll)(a).size()
- #define ppc __builtin_popcount
- #define ppcll __builtin_popcountll
- #define mem1(a) memset(a,-1,sizeof(a))
- #define mem0(a) memset(a,0,sizeof(a))
- #define endl "\n"
- #define lb lower_bound
- #define ub upper_bound
- template<class T> using ordered_set = tree<T, null_type,less<T>,rb_tree_tag,tree_order_statistics_node_update>;
- template<class T> using ordered_multiset = tree<T, null_type,less_equal<T>,rb_tree_tag,tree_order_statistics_node_update>;
- const ld PI = acos(-1.0);
- const int MOD = 1e9 +7;
- const ll INF = 1e18;
- // if(abs(a-b)<eps) --> if(a==b)
- // fixed << setprecision(n) -->printing decimal till n
- // hypot(a ,b) --> sqrt(a^2 + b^2)
- void solve(){
- ll n;
- cin>>n;
- vector<ll>C(n);
- for(int i=0;i<n;i++){
- cin>>C[i];
- }
- vector<ll>adj[n];
- for(int i=0;i<n-1;i++){
- ll a, b;
- cin>>a>>b;
- a--;
- b--;
- adj[a].pb(b);
- adj[b].pb(a);
- }
- ll time= 0;
- vector<ll>st(n) , en(n);
- function<void(ll ,ll)> dfs=[&](ll root , ll parent){
- st[root] = time;
- for(auto it : adj[root]){
- if(it==parent) continue;
- time++;
- dfs(it ,root);
- }
- en[root] = time;
- };
- dfs(0, -1);
- vector<ll>A(n);
- for(int i=0;i<n;i++){
- A[st[i]] = C[i];
- }
- //Sqrt Decomposition
- //two methods Online Query , Offline query
- //offline : we sort the queries and run in blocks
- //online : we precompute something nad combine blocks
- //Offline : --> (Sorting the queries ,MOs)
- //Coordinate compression so that we can apply add and remove function
- unordered_map<ll ,ll>mp1;
- ll ind =0;
- auto G = A;
- sort(all(G));
- G.resize(unique(all(G)) - G.begin());
- for(int i=0;i<sz(G);i++){
- mp1[G[i]] = ind++;
- }
- for(int i=0;i<n;i++){
- A[i] = mp1[A[i]];
- }
- vector<pair<pair<ll,ll> ,ll>>Queries;
- for(int i=0;i<n;i++){
- Queries.pb({{st[i] , en[i]} , i});
- }
- ll len = ceil(sqrt(n));
- auto comp=[&](pair<pair<ll, ll>, ll>a , pair<pair<ll, ll>, ll>b){
- ll s11 = (a.first.first/len);
- ll s12 = (b.first.first/len);
- ll s21 = (a.first.second);
- ll s22 = (b.first.second);
- if(s11==s12){
- return s11&1?s21>s22:s21<s22; //slight modification
- }
- return s11<s12;
- };
- sort(all(Queries) , comp);
- vector<ll>ans(n, 0);
- vector<ll>Cnt(2e5+1 , 0); //we remove and add by this
- ll l =0;
- ll r =-1;
- //l to r is the current range
- //We try to expand and shrink acc to new query
- ll distinct = 0;
- for(auto it : Queries){
- while(r < it.first.second){
- r++;
- Cnt[A[r]]++;
- if(Cnt[A[r]]==1) distinct++;
- }
- while(l > it.first.first){
- l--;
- Cnt[A[l]]++;
- if(Cnt[A[l]]==1) distinct++;
- }
- while(r > it.first.second){
- Cnt[A[r]]--;
- if(Cnt[A[r]]==0) distinct--;
- r--;
- }
- while(l < it.first.first){
- Cnt[A[l]]--;
- if(Cnt[A[l]]==0) distinct--;
- l++;
- }
- ans[it.second] = distinct;
- }
- for(auto it: ans){
- cout<<it<<" ";
- }
- cout<<endl;
- }
- int main(){
- IO;
- ll t =1;
- // cin>>t;
- for(int i=1;i<=t;i++){
- // cout<<"Case #"<<i<<": ";
- solve();
- }
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment