Submission #897219

#	Time	Username	Problem	Language	Result	Execution time	Memory
897219		ace5	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	791 ms	78276 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

using namespace std;

vector<int> segTree;
vector<vector<int>> g;
vector<int> path;
vector<int> dist;
map<int,int> uc;
vector<int> fs;
vector<int> ans;
vector<int> dp;
vector<int> c;

int v1 = -1,v2 = -1,n = -1;

void modify(int i,int x,int l,int r,int indV)
{
    if(i > r || i < l)
        return ;
    else if(l == r)
    {
        segTree[indV] += x;
        return ;
    }
    int m = (l+r)/2;
    modify(i,x,l,m,indV*2+1);
    modify(i,x,m+1,r,indV*2+2);
    segTree[indV] = segTree[indV*2+1] + segTree[indV*2+2];
    return ;
}
int get(int lx,int rx,int l,int r,int indV)
{
    if(l > rx || r < lx)
        return 0;
    else if(l >= lx && r <= rx)
    {
        return segTree[indV];
    }
    int m = (l+r)/2;
    int sl = get(lx,rx,l,m,indV*2+1);
    int sr = get(lx,rx,m+1,r,indV*2+2);
    return sl+sr;
}

pair<int,int> dfsf(int v,int p)
{
    int tv = v,td = 0;
    for(auto u:g[v])
    {
        if(u != p)
        {
            pair<int,int> ck = dfsf(u,v);
            if(ck.second+1 > td)
            {
                td = ck.second+1;
                tv = ck.first;
            }
        }
    }
    return {tv,td};
}
void dfsd(int v,int p,int d)
{
    dist[v] = d;
    for(auto u:g[v])
    {
        if(u != p)
        {
            dfsd(u,v,d+1);
        }
    }
    return ;
}
void dfsdp(int v,int p)
{
    dp[v] = 0;
    for(auto u:g[v])
    {
        if(u != p)
        {
            dfsdp(u,v);
            dp[v] = max(dp[v],dp[u]+1);

        }
    }
    return ;
}
void dfs(int v,int p,int b)
{
    //cout << v << ' ' << dp[v] << ' ' << fs[v] << ' ' << uc.size() << ' ' << (uc.size() != 0 ? (*uc.begin()).first : -1) << "\n";

    path.push_back(v);
    if(b == fs[v])
        ans[v] = uc.size();
    multiset<int> gl;
    for(auto u:g[v])
    {
        if(u != p)
        {
            gl.insert(-dp[u]);
        }
    }
    for(auto u:g[v])
    {
        if(u != p)
        {
            gl.erase(gl.find(-dp[u]));
            if(gl.size() != 0)
            {
                int mg = -(*gl.begin());
                modify(max(0,int(int(path.size())-mg-2)),1,0,n-1,0);
                modify(int(path.size())-1,-1,0,n-1,0);
                int fu = int(path.size())-dp[u]-1,su = int(path.size())-dp[u]-2;
               /* if(u == 1)
                {
                    cout << mg << ' ' << fu << ' ' << su << "\n";
                }*/
                if(fu >= 0 && get(0,fu,0,n-1,0) == 0)
                {
                    uc[c[path[fu]]]++;
                }
                if(su >= 0 && get(0,su,0,n-1,0) == 0)
                {
                    uc[c[path[su]]]++;
                }
                dfs(u,v,b);
                if(fu >= 0 && get(0,fu,0,n-1,0) == 0)
                {
                    uc[c[path[fu]]]--;
                    if(!uc[c[path[fu]]])
                        uc.erase(c[path[fu]]);
                }
                if(su >= 0 && get(0,su,0,n-1,0) == 0)
                {
                    uc[c[path[su]]]--;
                    if(!uc[c[path[su]]])
                        uc.erase(c[path[su]]);
                }
                modify(max(0,int(int(path.size())-mg-2)),-1,0,n-1,0);
                modify(int(path.size())-1,1,0,n-1,0);
            }
            else
            {
                int fu = int(path.size())-dp[u]-1,su = int(path.size())-dp[u]-2;
                if(fu >= 0 && get(0,fu,0,n-1,0) == 0)
                {
                    uc[c[path[fu]]]++;
                }
                if(su >= 0 && get(0,su,0,n-1,0) == 0)
                {
                    uc[c[path[su]]]++;
                }
                dfs(u,v,b);
                if(fu >= 0 && get(0,fu,0,n-1,0) == 0)
                {
                    uc[c[path[fu]]]--;
                    if(!uc[c[path[fu]]])
                        uc.erase(c[path[fu]]);
                }
                if(su >= 0 && get(0,su,0,n-1,0) == 0)
                {
                    uc[c[path[su]]]--;
                    if(!uc[c[path[su]]])
                        uc.erase(c[path[su]]);
                }
            }

            gl.insert(-dp[u]);
        }
    }
    path.pop_back();
    return ;
}

signed main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    int m;
    cin >> n >> m;
    segTree.resize(4*n);
    g.resize(n);
    dist.resize(n);
    fs.resize(n);
    ans.resize(n);
    dp.resize(n);
    c.resize(n);
    for(int i = 0;i < n-1;++i)
    {
        int a,b;
        cin >> a >> b;
        a--;
        b--;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    //cout << 1;
    for(int i = 0;i < n;++i)
        cin >> c[i];
    v1 = dfsf(0,-1).first;
    v2 = dfsf(v1,-1).first;
    dfsd(v1,-1,0);
    vector<int> d2 = dist;
    dfsd(v2,-1,0);
    //cout << 1;
    for(int j =0 ;j < n;++j)
    {
        if(d2[j] > dist[j])
        {
            fs[j] = 1;
        }
        else
            fs[j] = 2;
    }
   // cout << v1 << ' ' << v2 << ' ';
    //cout << 1;
    dfsdp(v1,-1);
    dfs(v1,-1,1);
    dfsdp(v2,-1);
    dfs(v2,-1,2);
   // cout << 1;
    for(int j = 0;j < n;++j)
    {
        cout << ans[j] << "\n";
    }
    return 0;
}

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