Submission #100928

#	Time	Username	Problem	Language	Result	Execution time	Memory
100928		Alexa2001	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	1197 ms	62584 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

using namespace std;


const int Nmax = 2e5 + 5;

int down[Nmax], ans[Nmax], level[Nmax], st[Nmax], max_ex[Nmax];
bool ap[Nmax];
vector<int> v[Nmax], collect_ans[Nmax], edge[Nmax];
int n, M, nr, Distincte;
int tip[Nmax], S[Nmax], tip_now[Nmax];


void read()
{
    int i, x, y;
    cin >> n >> M;
    for(i=1; i<n; ++i)
    {
        cin >> x >> y;
        v[x].push_back(y);
        v[y].push_back(x);
    }

    for(i=1; i<=n; ++i) cin >> tip[i];
}

void print_sol()
{
    int i;
    for(i=1; i<=n; ++i) cout << ans[i] << '\n';
}

void max_to(int &x, int y) { if(x < y) x = y; }


void dfs0(int node, int dad = 0)
{
    level[node] = level[dad] + 1;
    for(auto it : v[node])
        if(it != dad)
            dfs0(it, node);
}

pair<int,int> diametru()
{
    int i, A = 1, B = 1;

    dfs0(1);
    for(i=1; i<=n; ++i)
        if(level[i] > level[A]) A = i;

    dfs0(A);
    for(i=1; i<=n; ++i)
        if(level[i] > level[B]) B = i;

    return {A, B};
}

void dfs(int node, int dad = 0)
{
    down[node] = 1;
    level[node] = level[dad] + 1;

    for(auto it : v[node])
        if(it != dad)
        {
            dfs(it, node);
            down[node] = max(down[node], down[it] + 1);
        }

    int i, mx = 0;
    for(i=0; i<v[node].size(); ++i)
        if(v[node][i] != dad)
        {
            max_ex[v[node][i]] = mx;
            max_to(mx, down[v[node][i]]);
        }

    mx = 0;
    for(i=v[node].size()-1; i>=0; --i)
        if(v[node][i] != dad)
        {
            max_to(max_ex[v[node][i]], mx);
            max_to(mx, down[v[node][i]]);
        }
}

int bs(int val)
{
    int st = 1, dr = nr, mid;

    while(st <= dr)
    {
        mid = (st + dr) / 2;
        if(level[S[mid]] <= val) st = mid + 1;
            else dr = mid - 1;
    }
    return dr;
}

void solve(int node, int dad = 0)
{
    collect_ans[S[bs(level[node] - down[node] + 1)]].push_back(node);
    tip_now[node] = tip[dad];

    for(auto it : v[node])
        if(it != dad)
        {
            int pos, val, lnr;
            pos = bs(level[node] - max_ex[it]);
            lnr = nr;

            edge[S[pos]].push_back(it);

            val = S[pos+1];
            S[pos+1] = it;
            nr = pos+1;

            solve(it, node);

            S[pos+1] = val;
            nr = lnr;
        }
}

void extract_ans(int node)
{
    bool ok = 0;

    if(!ap[tip_now[node]])
    {
        ok = 1;
        ap[tip_now[node]] = 1;
        ++Distincte;
    }

    for(auto it : collect_ans[node])
        max_to(ans[it], Distincte);

    for(auto it : edge[node])
        extract_ans(it);

    if(ok)
    {
        ap[tip_now[node]] = 0;
        --Distincte;
    }
}

void Solve(int node)
{
    dfs(node);
    int i;
    for(i=0; i<=nr; ++i) S[i] = 0;
    nr = 0;
    for(i=0; i<=n; ++i) edge[i].clear(), collect_ans[i].clear(), ap[i] = 0;
    Distincte = 0;

    solve(node);
    for(auto it : edge[0]) extract_ans(it);
}

int main()
{
//    freopen("input", "r", stdin);
    cin.sync_with_stdio(false);

    read();
    auto P = diametru();

    Solve(P.first);
    Solve(P.second);

    print_sol();

    return 0;
}

Compilation message (stderr)

joi2019_ho_t5.cpp: In function 'void dfs(int, int)':
joi2019_ho_t5.cpp:74:15: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for(i=0; i<v[node].size(); ++i)
              ~^~~~~~~~~~~~~~~

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