Submission #201094

#TimeUsernameProblemLanguageResultExecution timeMemory
201094EntityITConstruction of Highway (JOI18_construction)C++14
100 / 100
376 ms18800 KiB
#include<bits/stdc++.h>

using namespace std;

#define all(x) (x).begin(), (x).end()
#define sz(x) ( (int)(x).size() )
using LL = long long;

mt19937 rng( (uint32_t)chrono::steady_clock::now().time_since_epoch().count() );

int n;
vector<int> c, par;
vector<vector<int> > gr;

vector<int> nChild;
void calNChild(int u) {
  nChild[u] = 1;
  for (int v : gr[u]) calNChild(v), nChild[u] += nChild[v];
}

vector<int> be, leader;
void hld(int u) {
  static int iTime = 0;

  be[u] = iTime++;
  if (leader[u] == -1) leader[u] = u;

  int big = -1;
  for (int v : gr[u]) if (big == -1 || nChild[big] < nChild[v]) big = v;

  if (big != -1) leader[big] = leader[u], hld(big);
  for (int v : gr[u]) if (v != big) hld(v);
}

struct Bit {
  vector<int> a;
  Bit(int nNode = 0) { a.assign(nNode, 0); }
  void reset(int nNode = 0) { a.assign(nNode, 0); }

  void upd(int pos, int val) {
    for (int i = pos; i < sz(a); i |= i + 1) a[i] += val;
  }
  int get(int pos) {
    int ret = 0;
    for (int i = pos; i >= 0; i = (i & (i + 1) ) - 1) ret += a[i];
    return ret;
  }
} cnt;

struct Chain { int l, r, val; };
vector<vector<Chain> > chain;

LL get(int u) {
  LL ret = 0;

  vector<Chain> tmp;
  for (int lead; ~u; u = par[lead]) {
    lead = leader[u];

    int r = be[u], prvSz = sz(tmp);
    for (int i = sz(chain[lead]) - 1; i >= 0; --i) {
      if (chain[lead][i].l > r) break ;
      tmp.emplace_back(Chain{ chain[lead][i].l, min(r, chain[lead][i].r), chain[lead][i].val } );
    }
    reverse(tmp.begin() + prvSz, tmp.end() );
  }

  for (auto _ : tmp) {
    ret += (LL)(_.r - _.l + 1) * cnt.get(_.val - 1);
    cnt.upd(_.val, _.r - _.l + 1);
  }
  for (auto _ : tmp) cnt.upd(_.val, -(_.r - _.l + 1) );

  return ret;
}

void upd(int u, int val) {
  for (int lead; ~u; u = par[lead]) {
    lead = leader[u];

    int l = be[lead], r = be[u];
    while (sz(chain[lead]) && chain[lead].back().r <= r) chain[lead].pop_back();
    if (sz(chain[lead]) ) chain[lead].back().l = r + 1;
    chain[lead].emplace_back(Chain{ l, r, val } );
  }
}

int main() {
  ios_base::sync_with_stdio(0); cin.tie(0);

  #ifdef FourLeafClover
  freopen("input", "r", stdin);
  #endif // FourLeafCLover

  cin >> n;

  c.assign(n, 0);
  vector<int> com;
  for (auto &i : c) {
    cin >> i;
    com.emplace_back(i);
  }
  sort(all(com) ); com.erase(unique(all(com) ), com.end() );
  for (auto &i : c) i = (int)(lower_bound(all(com), i) - com.begin() );

  par.assign(n, -1);
  gr.assign(n, {} );
  vector<pair<int, int> > edge;
  for (int i = 1; i < n; ++i) {
    int u, v; cin >> u >> v; --u; --v;
    edge.emplace_back(u, v);
    gr[u].emplace_back(v);
    par[v] = u;
  }

  nChild.assign(n, 0);
  calNChild(0);

  be.assign(n, -1);
  leader.assign(n, -1);
  hld(0);

  cnt.reset(sz(com) );

  chain.assign(n, {} );
  chain[0].emplace_back(Chain{ 0, 0, c[0] } );

  for (auto &e : edge) {
    int u = e.first, v = e.second;

    cout << get(u) << '\n';

    upd(v, c[v]);
  }

  return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...