Submission #1135598

#TimeUsernameProblemLanguageResultExecution timeMemory
1135598Hamed_GhaffariConstruction of Highway (JOI18_construction)C++20
100 / 100
198 ms22328 KiB
#include<bits/stdc++.h>
using namespace std;
#pragma GCC optimize("O3,unroll-loops")
#pragma GCC target("avx2,bmi,bmi2,sse4,sse4.2,lzcnt,popcnt")

using ll = long long;
using pii = pair<int, int>;

#define X first
#define Y second
#define all(x) x.begin(), x.end()
#define pb push_back
#define SZ(x) int(x.size())

const int MXN = 1e5+5;

namespace Fen {
  int fen[MXN];
  vector<pii> his;
  void updx(int p, int x, bool save=1) {
    if(save) his.pb(pii(p, x));
    for(p+=2; p<MXN; p+=p&-p) fen[p] += x;
  }
  int getx(int p) {
    int res=0;
    for(p+=2; p; p-=p&-p) res += fen[p];
    return res;
  }
  void clear() {
    while(SZ(his)) {
      updx(his.back().X, -his.back().Y, 0);
      his.pop_back();
    }
  }
}

int n, A[MXN], B[MXN], C[MXN];
vector<int> g[MXN];
int sz[MXN], hd[MXN], par[MXN], stt[MXN], fnt[MXN], timer;
vector<pii> st[MXN];

void compress() {
  vector<int> cmp;
  for(int i=1; i<=n; i++) cmp.pb(C[i]);
  sort(all(cmp));
  cmp.resize(unique(all(cmp))-cmp.begin());
  for(int i=1; i<=n; i++)
    C[i] = lower_bound(all(cmp), C[i])-cmp.begin();
}

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

void dfs2(int v) {
  stt[v] = timer++;
  if(!hd[v]) hd[v] = v;
  int big=-1;
  for(int u : g[v])
    if(u!=par[v] && (big==-1 || sz[u]>sz[big])) big = u;
  if(big != -1) {
    hd[big] = hd[v];
    dfs2(big);
    fnt[v] = fnt[big];
  }
  else {
    fnt[v] = timer-1;
  }
  for(int u : g[v])
    if(u!=par[v] && u!=big)
      dfs2(u);
  if(st[hd[v]].empty() || st[hd[v]].back().Y!=C[v])
    st[hd[v]].pb(pii(stt[v], C[v]));
  else st[hd[v]].back().X = stt[v];
}

ll upd(int v, int x) {
  vector<pii> tmp;
  ll res=0;
  while(v) {
    while(!st[hd[v]].empty() && st[hd[v]].back().X<=stt[v]) {
      if(SZ(st[hd[v]])==1 || st[hd[v]][SZ(st[hd[v]])-2].X>stt[v]) {
        tmp.pb(pii(st[hd[v]].back().Y, stt[v]-st[hd[v]].back().X+1));
        st[hd[v]].pop_back();
        if(!st[hd[v]].empty() && st[hd[v]][SZ(st[hd[v]])-2].X!=stt[v]+1) st[hd[v]].pb(pii(stt[v]+1, tmp.back().X));
        else if(stt[v]+1<=fnt[v]) st[hd[v]].pb(pii(stt[v]+1, tmp.back().X));
      }
      else {
        tmp.pb(pii(st[hd[v]].back().Y, st[hd[v]][SZ(st[hd[v]])-2].X-st[hd[v]].back().X));
        st[hd[v]].pop_back();
      }
    }
    st[hd[v]].pb(pii(stt[hd[v]], x));
    while(SZ(tmp)) {
      res += 1ll*tmp.back().Y*Fen::getx(tmp.back().X-1);
      Fen::updx(tmp.back().X, tmp.back().Y);
      tmp.pop_back();
    }
    v = par[hd[v]];
  }
  Fen::clear();
  return res;
}

int32_t main() {
  cin.tie(0); cout.tie(0); ios_base::sync_with_stdio(0);
  cin >> n;
  for(int i=1; i<=n; i++) cin >> C[i];
  compress();
  for(int i=0; i<n-1; i++) {
    cin >> A[i] >> B[i];
    g[A[i]].pb(B[i]);
    g[B[i]].pb(A[i]);
  }
  dfs1(1);
  dfs2(1);
  for(int i=0; i<n-1; i++)
    cout << upd(A[i], C[B[i]]) << '\n';
  return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...