Submission #109991

#	Time	Username	Problem	Language	Result	Execution time	Memory
109991		IOrtroiii	Construction of Highway (JOI18_construction)	C++14	100 / 100	673 ms	104508 KiB

construction

#include <bits/stdc++.h>

using namespace std;

const int N = 100100;

struct Data {
   int from, to, color;
   Data(int from = 0,int to = 0,int color = 0) : from(from), to(to), color(color) {}
};

int n;
vector<int> g[N];
int from[N], to[N], color[N];
int child[N], nxt[N], root[N];
int par[N], pos[N], ft[N];
deque<Data> dq[N];

void add(int p,int x) { for (; p <= n; p += p & -p) ft[p] += x; }
int get(int p) { int ans = 0; for (; p > 0; p -= p & -p) ans += ft[p]; return ans; }

void dfs(int u) {
   child[u] = 1;
   for (int v : g[u]) {
      par[v] = u;
      dfs(v);
      child[u] += child[v];
   }
}

void hld(int u,int root) {
   ::root[u] = root;
   if (u == root) pos[u] = 1;
   int nxt = 0;
   for (int v : g[u]) {
      if (child[v] > child[nxt]) nxt = v;
   }
   if (nxt) {
      pos[nxt] = pos[u] + 1;
      hld(nxt, root);
   }
   for (int v : g[u]) if (v != nxt) {
      hld(v, v);
   }
}

void get_hld(int u) {
   long long ans = 0;
   vector< pair<int, int> > change;
   while (u) {
      int p = pos[u];
      u = root[u];
      int ptr = 0;
      while (ptr < dq[u].size() && dq[u][ptr].from <= p) ptr++;
      ptr--;
      for (int i = ptr; i >= 0; --i) {
         int cnt = min(p, dq[u][i].to) - dq[u][i].from + 1;
         ans += (long long) cnt * get(dq[u][i].color - 1);
         add(dq[u][i].color, cnt);
         change.push_back({dq[u][i].color, cnt});
      }
      u = par[u];
   }

   for (auto p : change) add(p.first, -p.second);

   printf("%lld\n", ans);
}

void upd_hld(int u,int color) {
   while (u) {
      int p = pos[u];
      u = root[u];
      while (dq[u].size()) {
         if (dq[u][0].to <= p) dq[u].pop_front();
         else { dq[u][0].from = p + 1;  break; }
      }
      dq[u].push_front(Data(1, p, color));
      u = par[u];
   }
}

int main() {
   scanf("%d", &n);
   for (int i = 1; i <= n; ++i) {
      scanf("%d", color + i);
   }
   {
      vector<int> comp;
      for (int i = 1; i <= n; ++i) comp.push_back(color[i]);
      sort(comp.begin(), comp.end());
      comp.resize(unique(comp.begin(), comp.end()) - comp.begin());
      for (int i = 1; i <= n; ++i) color[i] = lower_bound(comp.begin(), comp.end(), color[i]) - comp.begin() + 1;
   }
   for (int i = 1; i < n; ++i) {
      scanf("%d %d", from + i, to + i);
      g[from[i]].push_back(to[i]);
   }
   dfs(1);
   hld(1, 1);
   for (int i = 1; i < n; ++i) {
      get_hld(from[i]);
      upd_hld(to[i], color[to[i]]);
   }
}

Compilation message (stderr)

construction.cpp: In function 'void get_hld(int)':
construction.cpp:54:18: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
       while (ptr < dq[u].size() && dq[u][ptr].from <= p) ptr++;
              ~~~~^~~~~~~~~~~~~~
construction.cpp: In function 'int main()':
construction.cpp:84:9: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
    scanf("%d", &n);
    ~~~~~^~~~~~~~~~
construction.cpp:86:12: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
       scanf("%d", color + i);
       ~~~~~^~~~~~~~~~~~~~~~~
construction.cpp:96:12: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
       scanf("%d %d", from + i, to + i);
       ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~

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