Submission #330170

#	Time	Username	Problem	Language	Result	Execution time	Memory
330170		VROOM_VARUN	Construction of Highway (JOI18_construction)	C++14	100 / 100	437 ms	27488 KiB

construction

/*
ID: varunra2
LANG: C++
TASK: highways
*/

#include <bits/stdc++.h>
using namespace std;

#ifdef DEBUG
#include "lib/debug.h"
#define debug(...) cerr << "[" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#define debug_arr(...) \
  cerr << "[" << #__VA_ARGS__ << "]:", debug_arr(__VA_ARGS__)
#pragma GCC diagnostic ignored "-Wsign-compare"
//#pragma GCC diagnostic ignored "-Wunused-parameter"
//#pragma GCC diagnostic ignored "-Wunused-variable"
#else
#define debug(...) 42
#endif

#define EPS 1e-9
#define IN(A, B, C) assert(B <= A && A <= C)
#define INF (int)1e9
#define MEM(a, b) memset(a, (b), sizeof(a))
#define MOD 1000000007
#define MP make_pair
#define PB push_back
#define all(cont) cont.begin(), cont.end()
#define rall(cont) cont.end(), cont.begin()
#define x first
#define y second

const double PI = acos(-1.0);
typedef long long ll;
typedef long double ld;
typedef pair<int, int> PII;
typedef map<int, int> MPII;
typedef multiset<int> MSETI;
typedef set<int> SETI;
typedef set<string> SETS;
typedef vector<int> VI;
typedef vector<PII> VII;
typedef vector<VI> VVI;
typedef vector<string> VS;

#define rep(i, a, b) for (int i = a; i < (b); ++i)
#define trav(a, x) for (auto& a : x)
#define sz(x) (int)(x).size()
typedef pair<int, int> pii;
typedef vector<int> vi;
#pragma GCC diagnostic ignored "-Wsign-compare"
// util functions
int n;
VVI adj;
VI wch;   // which chain
VI pos;   // which position in the chain
VI sub;   // subtree size
VI dep;   // distance from the root
VI live;  // liveliness
VII edgs;

void init() {
  adj.resize(n);
  wch.resize(n);
  pos.resize(n);
  sub.resize(n);
  dep.resize(n);
  live.resize(n);
  edgs.resize(n - 1);
}

struct chain {
  VII vals;
  int ind;
  int siz;
  int par;

  void init(int _ind, int _par) {
    ind = _ind;
    par = _par;
    siz = 0;
  }

  void ins(int node) {
    if (vals.empty())
      vals.PB(MP(live[node], 1));
    else if (vals.back().x == live[node])
      vals.back().y++;
    else
      vals.PB(MP(live[node], 1));
    siz++;
    wch[node] = ind;
    pos[node] = siz;
  }

  void done() {
    // call this when you are done with this chain
    reverse(all(vals));
  }

  VII get(int st, int val) {
    // get from 0 -> st
    VII ret;
    int sum = 0;
    while (!vals.empty()) {
      auto x = vals.back();
      int sum1 = sum + x.y;
      if (st <= sum1) {
        ret.PB(MP(x.x, st - sum));
        vals.back().y -= (st - sum);
        if (vals.back().y == 0) vals.pop_back();
        break;
      } else
        ret.PB(x);
      sum = sum1;
      vals.pop_back();
    }
    vals.PB(MP(val, st));
    return ret;
  }

  void deb() {
    debug("new chain");
    debug(ind);
    debug(par);
    debug(siz);
    debug(vals);
    debug("done");
  }
};

vector<chain> chains;

struct BIT {
  VI fen;
  int n;

  void init(int _n) {
    n = _n;
    fen.assign(n + 1, 0);
  }

  void upd(int ind, int val) {
    ind++;
    while (ind <= n) {
      fen[ind] += val;
      ind += (ind & -ind);
    }
  }

  int qry(int ind) {
    int ret = 0;
    ind++;
    while (ind >= 1) {
      ret += fen[ind];
      ind -= (ind & -ind);
    }
    return ret;
  }

  int qry(int l, int r) { return qry(r) - qry(l - 1); }
};

MPII cmp;

// first is val, second is count

void cmprs(VII vals) {
  cmp.clear();
  sort(all(vals));
  vals.resize(unique(all(vals)) - vals.begin());
  int cnt = 0;
  trav(x, vals) cmp[x.x] = cnt++;
}

ll calcinvs(VII vals) {
  ll ret = 0;
  BIT bit;
  reverse(all(vals));
  bit.init(sz(vals));
  cmprs(vals);
  for (int i = 0; i < sz(vals); i++) {
    int cur = bit.qry(cmp[vals[i].x] - 1);
    ret += cur * vals[i].y;
    PII ud = MP(cmp[vals[i].x], vals[i].y);
    bit.upd(ud.x, ud.y);
  }
  return ret;
}

VII gen(int node, int og = -1) {
  if (og == -1) og = node;
  VII ret;
  if (node == -1) return ret;
  ret = gen(chains[wch[node]].par, og);
  auto x = chains[wch[node]].get(pos[node], live[og]);
  trav(xx, x) ret.PB(xx);
  return ret;
}

ll qry(int node, int og) { return calcinvs(gen(node, og)); }

void dfs1(int u, int v) {
  sub[u] = 1;
  for (auto& x : adj[u]) {
    if (x == v) continue;
    dep[x] = dep[u] + 1;
    dfs1(x, u);
    sub[u] += sub[x];
  }
}

void dfs2(int u, int v) {
  if (sub[u] == 1) return;
  if (u == 0) {
    chain chn;
    chn.init(0, -1);
    chn.ins(0);
    chains.PB(chn);
  }
  int mx = adj[u][0];
  if (v == mx) mx = adj[u][1];
  trav(x, adj[u]) {
    if (x == v) continue;
    if (sub[mx] < sub[x]) mx = x;
  }
  trav(x, adj[u]) {
    if (x == v) continue;
    if (x == mx) {
      chains[wch[u]].ins(x);
    } else {
      chain nw;
      nw.init(sz(chains), u);
      nw.ins(x);
      chains.PB(nw);
    }
  }
  trav(x, adj[u]) {
    if (x == v) continue;
    dfs2(x, u);
  }
}

int main() {
  // #ifndef ONLINE_JUDGE
  // freopen("highways.in", "r", stdin);
  // freopen("highways.out", "w", stdout);
  // #endif
  cin.sync_with_stdio(0);
  cin.tie(0);

  cin >> n;

  init();

  for (int i = 0; i < n; i++) {
    cin >> live[i];
  }

  for (int i = 0; i < n - 1; i++) {
    int u, v;
    cin >> u >> v;
    u--, v--;
    edgs[i] = MP(u, v);
    adj[u].PB(v);
    adj[v].PB(u);
  }

  dfs1(0, -1);
  dfs2(0, -1);

  trav(x, chains) x.done();

  // debug(wch);
  // debug(pos);

  // trav(x, chains) x.deb();

  for (int i = 0; i < n - 1; i++) {
    cout << qry(edgs[i].x, edgs[i].y) << '\n';
  }

  // trav(x, chains) x.deb();

  return 0;
}

Compilation message (stderr)

construction.cpp: In member function 'void chain::deb()':
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
   19 | #define debug(...) 42
      |                    ^~
construction.cpp:124:5: note: in expansion of macro 'debug'
  124 |     debug("new chain");
      |     ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
   19 | #define debug(...) 42
      |                    ^~
construction.cpp:125:5: note: in expansion of macro 'debug'
  125 |     debug(ind);
      |     ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
   19 | #define debug(...) 42
      |                    ^~
construction.cpp:126:5: note: in expansion of macro 'debug'
  126 |     debug(par);
      |     ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
   19 | #define debug(...) 42
      |                    ^~
construction.cpp:127:5: note: in expansion of macro 'debug'
  127 |     debug(siz);
      |     ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
   19 | #define debug(...) 42
      |                    ^~
construction.cpp:128:5: note: in expansion of macro 'debug'
  128 |     debug(vals);
      |     ^~~~~
construction.cpp:19:20: warning: statement has no effect [-Wunused-value]
   19 | #define debug(...) 42
      |                    ^~
construction.cpp:129:5: note: in expansion of macro 'debug'
  129 |     debug("done");
      |     ^~~~~

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