Submission #314118

#TimeUsernameProblemLanguageResultExecution timeMemory
314118davooddkareshkiConstruction of Highway (JOI18_construction)C++17
100 / 100
799 ms36080 KiB
#include<bits/stdc++.h>

using namespace std;

typedef long long ll;
#define int long long
#define F first
#define S second
#define pii pair<int,int>
#define mpr make_pair

const int maxn = 1e5+10;
const int mod = 1e9+7;
const ll inf = 1e18+10;

int c[maxn], n;
int U[maxn], V[maxn];
int st[maxn], par[maxn], dis[maxn], hp[maxn];
vector<int> g[maxn];

struct HLD
{
    vector<int> comp;
    set<pii> se;

    vector<pii> upd(int pos, int val)
    {
        vector<pii> res;
        auto it_end = se.upper_bound(mpr(pos,inf));
        vector<pii> del;
        auto it = se.begin();

        /*if(pos == 3 && val == 3)
        {
            for(auto x : se)
                cout<< x.F <<" "<< x.S <<"\n";
            cout<<"\n";
        }*/

        for(it = se.begin(); it != it_end; it++)
        {
            del.push_back(*it);
            if(it == se.begin())
                res.push_back(*it);
            else
            {
                auto it2 = it; it2--;
                int num = (*it).F - (*it2).F;
                res.push_back({num,(*it).S});
            }
        }
        if(it == se.begin())
            res.push_back({pos,(*it).S});
        else
        {
            if(it != se.end())
            {
                auto it2 = it; it2--;
                res.push_back({pos-(*it2).F, (*it).S});
            }
        }

        for(auto X : del) se.erase(X);
        se.insert({pos,val});

        return res;
    }
} hld[maxn];

void pfs(int v)
{
    st[v] = 1;
    for(auto u : g[v])
        if(u != par[v])
        {
            par[u] = v;
            pfs(u);
            st[v] += st[u];
        }
}

void dfs(int v)
{
    hld[hp[v]].comp.push_back(v);

    for(auto u : g[v])
        if(u != par[v])
        {
            if(st[v] > 2 * st[u])
            {
                hp[u] = u;
                dis[u] = 1;
            }
            else
            {
                hp[u] = hp[v];
                dis[u] = dis[v] + 1;
            }

            dfs(u);
        }
}

struct fenwick
{
    int fen[maxn];
    fenwick() {memset(fen, 0, sizeof fen);}

    void add(int pos, int val)
    {
        for(; pos <= n; pos |= (pos+1))
            fen[pos] += val;
    }

    int qu(int pos)
    {
        int res = 0;
        for(; pos >= 0; pos = (pos&(pos+1))-1)
            res += fen[pos];
        return res;
    }
} fen;

int upd(int v, int val)
{
    vector<pii> Y;
    while(v != 0)
    {
        vector<pii> X = hld[hp[v]].upd(dis[v],val);
        reverse(X.begin(), X.end());
        for(auto x : X) Y.push_back(x);
        v = par[hp[v]];
    }

    int ans = 0;
    //cout<<"\n";
    for(auto X : Y)
    {
        int num = X.F, val = X.S;
        ans += fen.qu(val-1) * num;
        fen.add(val, num);
        //cout<< num <<" "<< val <<"\n";
    }

    for(auto X : Y)
    {
        int num = X.F, val = X.S;
        fen.add(val, -num);
    }

    return ans;
}

signed main()
{
    ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);

    cin>> n; vector<int> ve;
    for(int i = 1; i <= n; i++)
    {
        cin>> c[i];
        ve.push_back(c[i]);
    }
    sort(ve.begin(), ve.end());
    ve.resize(unique(ve.begin(), ve.end()) - ve.begin());
    for(int i = 1; i <= n; i++)
        c[i] = upper_bound(ve.begin(), ve.end(), c[i]) - ve.begin();

    for(int i = 1; i < n; i++)
    {
        cin>> U[i] >> V[i];
        g[U[i]].push_back(V[i]);
        g[V[i]].push_back(U[i]);
    }

    pfs(1);
    dis[1] = 1; hp[1] = 1;
    dfs(1);

    for(int i = 1; i <= n; i++)
    {
        if(hp[i] == i)
            hld[i].se.insert({(int)hld[i].comp.size(),n+1});
        //cout<< i <<" "<< hp[i] <<" "<< dis[i] <<" "<< c[i] <<"\n";
    }

    for(int i = 1, u, v; i < n; i++)
    {
        u = U[i]; v = V[i];
        cout<< upd(v,c[v]) <<"\n";
    }
}

/*
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 'int main()':
construction.cpp:187:20: warning: variable 'u' set but not used [-Wunused-but-set-variable]
  187 |     for(int i = 1, u, v; i < n; i++)
      |                    ^
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...