Submission #757831

#	Time	Username	Problem	Language	Result	Execution time	Memory
757831		taher	Transport (COCI19_transport)	C++17	130 / 130	419 ms	30500 KiB

Main

#include <bits/stdc++.h>

using namespace std;

const int N = 100005;

int n, poi, cnt, curCentroid;
int subSizes[N];
vector<array<long long, 2>> adj[N];
bool vis[N];
vector<long long> all, v[N], c[N], g[N];
long long ans = 0;
long long a[N];
long long fenw[2 * N];
int times[2 * N];

void modify(int x, long long add) {
  while (x < (int) all.size()) {
    if (times[x] != cnt) {
      times[x] = cnt;
      fenw[x] = 0;
    }
    fenw[x] += add;
    x |= (x + 1);
  }
}

long long get(int x) {
  long long sum = 0;
  while (x >= 0) {
    if (times[x] != cnt) {
      fenw[x] = 0;
    }
    sum += fenw[x];
    x = (x & (x + 1)) - 1;
  }
  return sum;
}

void initArrays() {
  for (int i = 0; i < n; i++) {
    vis[i] = false;
  }
  for (int i = 0; i < 2 * n; i++) {
    times[i] = -1;
  }
  return ;
}

int getSiz(int u, int prev) {
  subSizes[u] = 1;
  for (auto x : adj[u]) {
    if (!vis[x[0]] && x[0] != prev) {
      subSizes[u] += getSiz(x[0], u);
    }
  }
  return subSizes[u];
}

int locateCentroid(int u, int prev, int sz) {
  for (auto x : adj[u]) {
    if (vis[x[0]] || x[0] == prev) {
      continue;
    }
    if (subSizes[x[0]] > sz / 2) {
      return locateCentroid(x[0], u, sz);
    }
  }
  return u;
}

int centroid(int u, int prev) {
  int sz = getSiz(u, prev);
  return locateCentroid(u, prev, sz);
}

void dfs(int u, int prev, long long balance, long long minPrefixBalance, long long minSuffixBalance) {
  all.push_back(-minPrefixBalance + a[curCentroid]);
  v[poi].push_back(minPrefixBalance);
  minSuffixBalance += a[u];
  balance += a[u];
  all.push_back(balance);

  g[poi].push_back(minSuffixBalance);
  c[poi].push_back(balance);
  for (auto x : adj[u]) {
    if (x[0] == prev || vis[x[0]]) {
      continue;
    }
    balance += x[1];
    long long nMnPref = min(minPrefixBalance, balance);
    dfs(x[0], u, balance, nMnPref, min(0LL, minSuffixBalance) + x[1]);
    balance -= x[1];
  }
  return ;
}

void Decompose(int u, int prev) {
  u = centroid(u, prev);
  curCentroid = u;
  vis[u] = 1;
  poi = 0;
  cnt++;
  all.clear();
  for (auto x : adj[u]) {
    if (x[0] == prev || vis[x[0]]) {
      continue;
    }
    dfs(x[0], u, a[u] + x[1], a[u] + x[1], x[1]);
    ++poi;
  }
  // all.push_back(a[u]);
  sort(all.begin(), all.end());
  all.erase(unique(all.begin(), all.end()), all.end());

  auto Compress = [&](long long pts) {
    int x = lower_bound(all.begin(), all.end(), pts) - all.begin();
    return x;  
  };

  for (int i = 0; i < poi; i++) {
    for (int j = 0; j < (int) v[i].size(); j++) {
      long long minPrefixBalance = v[i][j] - a[u];
      if (minPrefixBalance + a[u] >= 0) {
        ans += 1;
      }
      if (minPrefixBalance < 0) {
        modify(Compress(-minPrefixBalance), +1);
      } else {
        modify(0, +1);
      }
    }
  }

  for (int i = 0; i < poi; i++) {
    for (int j = 0; j < (int) c[i].size(); j++) {
      long long minPrefixBalance = v[i][j] - a[u];
      if (minPrefixBalance < 0) {
        modify(Compress(-minPrefixBalance), -1);
      } else {
        modify(0, -1);
      }      
    }

    for (int j = 0; j < (int) c[i].size(); j++) {
      long long curBalance = c[i][j];
      long long minSuffixBalance = g[i][j];
      if (minSuffixBalance >= 0) {
        ans += 1;
        ans += get(Compress(curBalance));
      }
    }

    for (int j = 0; j < (int) c[i].size(); j++) {
      long long minPrefixBalance = v[i][j] - a[u];
      if (minPrefixBalance < 0) {
        modify(Compress(-minPrefixBalance), +1);
      } else {
        modify(0, +1);
      }
    }
  }

  for (int i = 0; i < poi; i++) {
    v[i].clear();
    c[i].clear();
    g[i].clear();
  }
  for (auto x : adj[u]) {
    if (!vis[x[0]] && x[0] != prev) {
      Decompose(x[0], u);
    }
  }
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cin >> n;
  initArrays();
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }
  for (int i = 0; i < n - 1; i++) {
    int f, t, cost;
    cin >> f >> t >> cost;
    --f, --t;
    adj[f].push_back({t, -cost});
    adj[t].push_back({f, -cost});
  }
  Decompose(0, -1);
  cout << ans << '\n';
  return 0;
}

• Subtask 1: 13 / 13
• Subtask 2: 13 / 13
• Subtask 3: 13 / 13
• Subtask 4: 13 / 13
• Subtask 5: 13 / 13
• Subtask 6: 13 / 13
• Subtask 7: 13 / 13
• Subtask 8: 13 / 13
• Subtask 9: 13 / 13
• Subtask 10: 13 / 13