Submission #1180722

#	Time	Username	Problem	Language	Result	Execution time	Memory
1180722		Zbyszek99	Unique Cities (JOI19_ho_t5)	C++20	100 / 100	689 ms	54072 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
#define ll long long
#define ld long double
#define ff first
#define ss second
#define vl vector<long long>
#define vi vector<int>
#define pii pair<int,int>
#define pll pair<long long, long long>
#define pb push_back
#define rep(i, b) for(int i = 0; i < (b); ++i)
#define rep2(i,a,b) for(int i = a; i <= (b); ++i)
#define rep3(i,a,b,c) for(int i = a; i <= (b); i+=c)
#define count_bits(x) __builtin_popcountll((x))
#define all(x) (x).begin(),(x).end()
#define siz(x) (int)(x).size()
#define forall(it,x) for(auto& it:(x))
using namespace std;
//mt19937 mt;void random(){mt.seed(chrono::time_point_cast<chrono::milliseconds>(chrono::high_resolution_clock::now()).time_since_epoch().count());}
//ll rand(ll a, ll b) {return a + (mt() % (b-a+1));}
const int INF = 1e9+50;
const ll INF_L = 1e18+40;
const ll MOD = 1e9+7;

const int tree_siz = 1024*512-1;
int drzewo[tree_siz+1];

void add_seg(int akt, int p1, int p2, int s1, int s2, int w)
{
    if(p2 < s1 || p1 > s2) return;
    if(p1 >= s1 && p2 <= s2)
    {
        drzewo[akt]+=w;
        return;
    }
    add_seg(akt*2,p1,(p1+p2)/2,s1,s2,w);
    add_seg(akt*2+1,(p1+p2)/2+1,p2,s1,s2,w);
}

int get_sum(int v)
{
    return drzewo[v] + (v != 1 ? get_sum(v/2) : 0);
}

int get_s(int v)
{
    return get_sum(tree_siz/2+1+v);
}

vi graph[200'001];
int dist[200'001];
int son_dist[200'001];
pii best;

void dfs_dist(int v, int pop, int d)
{
    dist[v] = d;
    best = max(best,{d,v});
    forall(it,graph[v])
    {
        if(it != pop)
        {
            dfs_dist(it,v,d+1);
        }
    }
}

int dfs_son(int v, int pop)
{   
    int d = 1;
    forall(it,graph[v])
    {
        if(it != pop)
        {
            d = max(d,dfs_son(it,v)+1);
        }
    }
    son_dist[v] = d;
    return d;
}

int ans[200'001];
int color[200'001];
set<int> cur_unq;
int zlicz[200'001];
vi cur_list;

void dfs_ans(int v, int pop)
{
    ans[v] = max(ans[v],siz(cur_unq));
    cur_list.pb(v);
    // cout << v << " dfs: ";
    // forall(it,cur_list) cout << it << " ";
    // cout << " list\n";
    // forall(it,cur_unq) cout << it << " ";
    // cout << " unq\n";
    // rep(d,siz(cur_list)) cout << get_s(d) << " ";
    // cout << "seg_tree\n\n";

    forall(it,graph[v])
    {
        if(it != pop)
        {
            int ins = -1;
            int ins2 = -1;
            int del_ind = dist[v]-son_dist[it];
            int del_ind2 = del_ind+1;
            if(dist[v] - 1 >= max(0,dist[v] - son_dist[v]+1)) add_seg(1,0,tree_siz/2,max(0,dist[v]-son_dist[v]+1),dist[v]-1,-1);
            forall(d,graph[it])
            {
                if(d != v)
                {
                    add_seg(1,0,tree_siz/2,max(0,dist[it] - son_dist[d]),dist[v],1);
                }
            }
            if(del_ind >= 0 && get_s(del_ind) == 0)
            {
         //       cout << cur_list[del_ind] << " " << color[cur_list[del_ind]] << " cur\n";
                ins = color[cur_list[del_ind]];
                zlicz[ins]++;
                cur_unq.insert(ins);
            }
            if(del_ind2 >= 0 && get_s(del_ind2) == 0)
            {
         //       cout << cur_list[del_ind] << " " << color[cur_list[del_ind]] << " cur\n";
                ins2 = color[cur_list[del_ind2]];
                zlicz[ins2]++;
                cur_unq.insert(ins2);
            }
          //  cout << ins << " " << del_ind << " " << it << " " << v << " add\n";
            dfs_ans(it,v);
            if(dist[v] - 1 >= max(0,dist[v] - son_dist[v]+1)) add_seg(1,0,tree_siz/2,max(0,dist[v]-son_dist[v]+1),dist[v]-1,1);
            forall(d,graph[it])
            {
                if(d != v)
                {
                    add_seg(1,0,tree_siz/2,max(0,dist[it] - son_dist[d]),dist[v],-1);
                }
            }
            if(ins != -1)
            {
                zlicz[ins]--;
                if(zlicz[ins] == 0)
                {
                    cur_unq.erase(cur_unq.find(ins));
                }
            }
            if(ins2 != -1)
            {
                zlicz[ins2]--;
                if(zlicz[ins2] == 0)
                {
                    cur_unq.erase(cur_unq.find(ins2));
                }
            }
        }
    }
    cur_list.pop_back();
}

void solve(int v)
{
  //  cout << "\nsolve\n";
    cur_unq = {};
    cur_list = {};
    rep2(i,1,tree_siz) drzewo[i] = 0;
    dfs_dist(v,v,0);
    dfs_son(v,v);
    dfs_ans(v,v);
}

int main()
{
    ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
    //random();
    int n,m;
    cin >> n >> m;
    rep(i,n-1)
    {
        int a,b;
        cin >> a >> b;
        graph[a].pb(b);
        graph[b].pb(a);
    }
    best = {0,0};
    dfs_dist(1,1,0);
    int v1 = best.ss;
    best = {0,0};
    dfs_dist(v1,v1,0);
    int v2 = best.ss;
  //  cout << v1 <<  " " << v2 <<  "v\n";
    rep2(i,1,n) cin >> color[i];
    solve(v1);
    solve(v2);
    rep2(i,1,n) cout << ans[i] << "\n";
}

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32