Submission #60307

#	Time	Username	Problem	Language	Result	Execution time	Memory
60307		XmtosX	Construction of Highway (JOI18_construction)	C++17	100 / 100	855 ms	89352 KiB

construction

#include <bits/stdc++.h>
using namespace std;
const int N=1e5+5;
int n,c[N],a[N],b[N],sz[N],par[N],nxt[N],lvl[N],bit[N];
vector<int> v[N];
vector<pair<int,int> > V[N];
void update (int x,int val)
{
    while (x<N)
    {
        bit[x]+=val;
        x+= (x&(-x));
    }
}
int query (int x)
{
    int sum=0;
    while (x)
    {
        sum+= bit[x];
        x-= (x&(-x));
    }
    return sum;
}
void dfs (int x,int p)
{
    par[x]=p;
    sz[x]=1;
    lvl[x]=lvl[p]+1;
    for (int i=0;i<v[x].size();i++)
    {
        dfs(v[x][i],x);
        sz[x]+=sz[v[x][i]];
        if (sz[v[x][i]]>sz[v[x][0]])
            swap(v[x][i],v[x][0]);
    }
}
void hld (int x)
{
    for (int i=0;i<v[x].size();i++)
    {
        if (i==0)
            nxt[v[x][i]]=nxt[x];
        else
            nxt[v[x][i]]=v[x][i];
        hld(v[x][i]);
    }
}
long long solve(int x)
{
    vector <pair<int,int> > p;
    vector <int> v;
    int z=x;
    while (x)
    {
        int l= (lvl[x]-lvl[nxt[x]]+1);
        x=nxt[x];
        vector <pair<int,int> > o;
        for (int i=V[x].size()-1;i>=0;i--)
        {
            int cnt=V[x][i].first;
            int y=V[x][i].second;

            if (cnt<l)
            {
                o.push_back(V[x][i]);
                v.push_back(y);
                l-=cnt;
            }
            else
            {
                o.push_back(make_pair(l,y));
                v.push_back(y);
                break;
            }
        }
        for (int i=o.size()-1;i>=0;i--)
            p.push_back(o[i]);
        x=par[x];
    }
    reverse(p.begin(),p.end());
    v.push_back(0);
    sort(v.begin(),v.end());
    v.resize(unique(v.begin(),v.end())-v.begin());
    for (int i=0;i<=v.size();i++)
        bit[i]=0;
    long long ans=0;
    int tot=0;
    for (int i=0;i<p.size();i++)
    {
        int y=p[i].second;
        long long cnt=p[i].first;
        x= lower_bound(v.begin(),v.end(),y)-v.begin();
        ans+= cnt*(tot-query(x));
        tot+=cnt;
        update(x,cnt);
    }
    return ans;
}
void change (int x)
{
    int q=c[x];
    while (x)
    {
        int l= lvl[x]-lvl[nxt[x]]+1;
        int y=l;
        x=nxt[x];
        while (!V[x].empty())
        {
            int i= (V[x].size()-1);
            int cnt=(V[x][i].first);
            if (cnt<=l)
            {
                l-=cnt;
                V[x].pop_back();
            }
            else
            {
                V[x][i].first-=l;
                break;
            }
        }
        if (!V[x].empty()&&V[x][V[x].size()-1].second==q)
            V[x][V[x].size()-1].first+=y;
        else
            V[x].push_back({y,q});
        x=par[x];
    }
}
int main()
{
    cin >>n;
    for (int i=1;i<=n;i++)
        cin >>c[i];
    for (int i=0;i<n-1;i++)
    {
        cin >>a[i]>>b[i];
        v[a[i]].push_back(b[i]);
    }
    dfs(1,0);
    nxt[1]=1;
    hld(1);
    V[1].push_back({1,c[1]});
    for (int i=0;i<n-1;i++)
    {
        cout <<solve(a[i])<<"\n";
        change(b[i]);
    }
    return 0;
}
/*
5
1 2 3 4 5
1 2
2 3
2 4
3 5
*/
/*
10
1 7 3 4 8 6 2 9 10 5
1 2
1 3
2 4
3 5
2 6
3 7
4 8
5 9
6 10
*/

Compilation message (stderr)

construction.cpp: In function 'void dfs(int, int)':
construction.cpp:30:19: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i=0;i<v[x].size();i++)
                  ~^~~~~~~~~~~~
construction.cpp: In function 'void hld(int)':
construction.cpp:40:19: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i=0;i<v[x].size();i++)
                  ~^~~~~~~~~~~~
construction.cpp: In function 'long long int solve(int)':
construction.cpp:85:19: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i=0;i<=v.size();i++)
                  ~^~~~~~~~~~
construction.cpp:89:19: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i=0;i<p.size();i++)
                  ~^~~~~~~~~
construction.cpp:53:9: warning: unused variable 'z' [-Wunused-variable]
     int z=x;
         ^

• Subtask 1: 7 / 7
• Subtask 2: 9 / 9
• Subtask 3: 84 / 84