Submission #465897

#	Time	Username	Problem	Language	Result	Execution time	Memory
465897		erray	Construction of Highway (JOI18_construction)	C++17	100 / 100	415 ms	21544 KiB

construction

// author: erray
#include <bits/stdc++.h>
 
using namespace std;
template<typename A, typename B> string to_string(const pair<A, B>& p);
template<typename A, typename B, typename C> string to_string(const tuple<A, B, C>& t);
template<typename A, typename B, typename C, typename D> string to_string(const tuple<A, B, C, D>& t);

string to_string(const string& s) {
  return '"' + s + '"';
}

string to_string(const char& c) {
  return string("'") + c + "'";
}

string to_string(const char *c) {
  return to_string(string(c));
}

string to_string(const bool& b) {
  return (b ? "true" : "false");
}

string to_string(const vector<bool>& v) {
  string res = "{";
  for (int i = 0; i < (int) v.size(); ++i) {
    if (i > 0) {
      res += ", ";
    }
    res += to_string(v[i]);
  }
  res += "}";
  return res;
}

template<size_t T> string to_string(const bitset<T>& bs) {
  return bs.to_string();
}

template<typename T> string to_string(queue<T> q) {
  string res = "{";
  size_t sz = q.size();
  while (sz--) {
    T cur = q.front();
    q.pop();
    if ((int) res.size() > 1) {
      res += ", ";
    }
    res += to_string(cur);
  }
  res += "}";
  return res;
}

template<typename T, class C> string to_string(priority_queue<T, vector<T>, C> pq) {
  string res = "{";
  while (!pq.empty()) {
    T cur = pq.top();
    pq.pop();
    if ((int) res.size() > 1) {
      res += ", ";
    }
    res += to_string(cur);
  }
  res += "}";
  return res;
}

template<typename T> string to_string(const T& v) {
  string res = "{";
  for (auto &el : v) {
    if ((int) res.size() > 1) {
      res += ", ";
    }
    res += to_string(el);
  }
  res += "}";
  return res;
}

template<typename A, typename B> string to_string(const pair<A, B>& p) {
  return '(' + to_string(p.first) + ", " + to_string(p.second) + ')';
}
template<typename A, typename B, typename C> string to_string(const tuple<A, B, C>& t) {
  return '(' + to_string(get<0>(t)) + ", " + to_string(get<1>(t)) + ", " + to_string(get<2>(t)) + ')';
}
template<typename A, typename B, typename C, typename D> string to_string(const tuple<A, B, C, D>& t) {
  return '(' + to_string(get<0>(t)) + ", " + to_string(get<1>(t)) + ", " + to_string(get<2>(t)) + ", " + to_string(get<3>(t)) + ')';
}

void debug_out(int size, bool first, string name) {
  vector<string> tmp = {name};
  vector<int> tmp2 = {size, first};
  cerr << endl;
}

constexpr int buffer_size = 255;

template<typename Head, typename... Tail> void debug_out(int size, bool first, string name, Head H, Tail... T) {
  string tmp;
  int off = 0;
  while ((!name.empty() && name[0] != ',') || off != 0) {
    tmp += name[0];
    name.erase(name.begin());
    char c = tmp.back();
    if (c == '{' || c == '(') {
      ++off;
    } else if (c == '}' || c == ')'){
      --off;
    }
  }
  if (!name.empty()) {
    name.erase(name.begin());
  }
  if (tmp[0] == ' ') {
    tmp.erase(tmp.begin());
  }

  string buff = to_string(H);
  if ((int) buff.size() + size + (int) tmp.size() > buffer_size - 5 && !first) {
    cerr << '\n';
    size = 0;
  }
  cerr << '[' << tmp << ": " << buff << "] ";
  debug_out(((int) buff.size() + size + (int) tmp.size() + 5) % buffer_size, false, name, T...);
}

#ifdef DEBUG
#define debug(...) cerr << "-> ", debug_out(3, true, string(#__VA_ARGS__), __VA_ARGS__)
#define here cerr << "-> " << __LINE__ << endl
#else
#define debug(...) void(37)
#define here void(37)
#endif

template<typename T> struct fenwick {
  int n;
  vector<T> tree;
  fenwick(int _n, T def = T{}) : n(_n) {
    tree.resize(n + 1);
    tree[0] = def;
  }

  void modify(int i, T x) {
    assert(i >= 0 && i < n);
    ++i;
    while (i <= n) {
      tree[i] += x;
      i += (i & -i);
    }
  }

  T get(int i) {
    assert(i >= 0 && i < n);
    ++i;
    T res = tree[0];
    while (i > 0) {
      res += tree[i];
      i -= (i & -i);
    }
    return res;
  }

  T get(int l, int r) {
    return get(r) - (l == 0 ? tree[0] : get(l - 1));
  }

  template<class F> int find_first(F op) {
    int cur = 0;
    T sum = tree[0];
    for (int i = __lg(n); i >= 0; --i) {
      if (cur + (1 << i) <= n && !op(sum + tree[cur + (1 << i)])) {
        cur += (1 << i);
        sum = sum + tree[cur];
      }
    }
    return (cur == n ? -1 : cur);
  }

  void clear() {
    tree.assign(n + 1, tree[0]);
  }
};
 
int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  int n;
  cin >> n;
  vector<int> c(n);
  for (int i = 0; i < n; ++i) {
    cin >> c[i];
  }

  vector<int> sc = c;
  sort(sc.begin(), sc.end());
  sc.erase(unique(sc.begin(), sc.end()), sc.end());
  int c_size = int(sc.size()) + 1;
  for (int i = 0; i < n; ++i) {
    c[i] = int(lower_bound(sc.begin(), sc.end(), c[i]) - sc.begin()) + 1;
  }
  vector<vector<int>> g(n);
  vector<array<int, 2>> es(n - 1);
  for (int i = 0; i < n - 1; ++i) {
    int x, y;
    cin >> x >> y;
    --x, --y;
    g[x].push_back(y);
    g[y].push_back(x);
    es[i] = {x, y};
  }
  
  vector<int> par(n, -1);
  vector<int> sz(n, 1);
  vector<int> ord;
  vector<int> d(n);
  function<void(int)> Dfs = [&](int v) {
    ord.push_back(v);
    for (auto& u : g[v]) {
      par[u] = v;
      g[u].erase(find(g[u].begin(), g[u].end(), v));
      d[u] = d[v] + 1;
      Dfs(u);
      sz[v] += sz[u];
      if (sz[u] > sz[g[v][0]]) {
        swap(u, g[v][0]);
      }    
    }
  };
  Dfs(0);
  debug(par, sz, ord, d);
  debug(c);
  vector<int> root(n);
  iota(root.begin(), root.end(), 0);
  for (auto v : ord) {
    if (!g[v].empty()) {
      root[g[v][0]] = root[v];
    }
  }

  debug(g, root);
  // number, size
  vector<vector<pair<int, int>>> ch(n);
  for (auto v : ord) {
    ch[root[v]].emplace_back(c[v], 1);
  }
  for (int i = 0; i < n; ++i) {
    reverse(ch[i].begin(), ch[i].end());
  }

  debug(ch);
  fenwick<int> inv(c_size);
  auto Apply = [&](int v) {
    vector<pair<int, int>> ns;
    int change = c[v];
    v = par[v];
    while (v != -1) {
      int size = d[v] - d[root[v]] + 1;
      v = root[v];
      int cur = 0;
      vector<pair<int, int>> pref;
      while (cur < size) {
        int dec = min(size - cur, ch[v].back().second);
        pref.emplace_back(ch[v].back().first, dec);
        if (dec == ch[v].back().second) {
          ch[v].pop_back();
        } else {
          ch[v].back().second -= dec;
        }
        cur += dec;
      }
      ns.insert(ns.end(), pref.rbegin(), pref.rend());
      ch[v].emplace_back(change, size);
      v = par[v];   
    }

    debug(ns);
    long long res = 0;
    for (auto[x, ct] : ns) {
      inv.modify(x, +ct);
      res += 1LL * ct * inv.get(x - 1);
    }
    for (auto[x, ct] : ns) {
      inv.modify(x, -ct);
    }
    return res;
  };

  for (auto[ignore, v] : es) {
    debug(v, c[v]);
    cout << Apply(v) << '\n';
    debug(ch);
  }
   
}

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