Submission #96623

#	Time	Username	Problem	Language	Result	Execution time	Memory
96623		markotee	Transport (COCI19_transport)	C++17	130 / 130	419 ms	13420 KiB

Main

#include <bits/stdc++.h>
#define all(x) begin(x), end(x)
#define dbg(x) cerr << #x << " = " << x << endl
#define _ << " " <<
using namespace std;
using ll = long long;

const int MXN = 100005;

int msz[MXN], siz[MXN];
int vis[MXN], fuel[MXN];

vector<ll> suff;
vector<ll> pref;
vector<int> nodes;
vector<pair<int, int>> adj[MXN];

void dfs(int x, int p) {
  siz[x] = 1;
  msz[x] = 0;
  nodes.push_back(x);
  for (auto &[y, w] : adj[x]) {
    if (y == p || vis[y]) continue;
    dfs(y, x);
    siz[x] += siz[y];
    msz[x] = max(msz[x], siz[y]);
  }
}

int centroid(int rt) {
  nodes.clear();
  dfs(rt, 0);
  int ret = 0, mn = 1e9;
  for (int x : nodes) {
    int mx = max(msz[x], siz[rt] - siz[x]);
    if (mn > mx) {
      mn = mx;
      ret = x;
    }
  }
  return ret;
}

void setuppref(int x, int p, ll sum, ll mn) {
  if (mn >= 0) pref.push_back(sum);
  for (auto &[y, w] : adj[x]) {
    if (y == p || vis[y]) continue;
    ll nxts = sum + fuel[y] - w;
    ll nxtm = min(mn, 0LL) + fuel[y] - w;
    setuppref(y, x, nxts, nxtm);
  }
}

ll setupsuff(int x, int p, ll sum, ll mx) {
  suff.push_back(mx);
  for (auto &[y, w] : adj[x]) {
    if (y == p || vis[y]) continue;
    ll nxt = sum + w - fuel[x];
    setupsuff(y, x, nxt, max(mx, nxt));
  }
}

ll solve() {
  ll ret = 0;
  sort(all(pref));
  sort(all(suff));
  for (int i = 0, j = 0; i < suff.size(); ++i) {
    while (j < pref.size() && suff[i] > pref[j]) {
      j++;
    }
    ret += pref.size() - j;
  }
  pref.clear();
  suff.clear();
  return ret;
}

int main() {
  ios::sync_with_stdio(0), cin.tie(0);
  int n;
  cin >> n;
  for (int i = 1; i <= n; ++i) {
    cin >> fuel[i];
  }
  for (int i = 2; i <= n; ++i) {
    int x, y, w;
    cin >> x >> y >> w;
    adj[x].push_back({y, w});
    adj[y].push_back({x, w});
  }
  ll sol = 0;
  queue<int> q;
  for (q.push(1); !q.empty(); q.pop()) {
    int c = centroid(q.front());
    setuppref(c, 0, 0, 0);
    setupsuff(c, 0, 0, 0);
    sol += solve();
    for (auto &[x, w] : adj[c]) {
      if (vis[x]) continue;
      setuppref(x, c, fuel[x] - w, min(fuel[x] - w, 0));
      setupsuff(x, c, w - fuel[c], max(w - fuel[c], 0));
      sol -= solve();
    }
    for (auto &[x, w] : adj[c]) {
      if (vis[x]) continue;
      q.push(x);
    }
    vis[c] = 1;
  }
  cout << sol - n << '\n';
}

Compilation message (stderr)

transport.cpp: In function 'void dfs(int, int)':
transport.cpp:22:19: warning: unused variable 'w' [-Wunused-variable]
   for (auto &[y, w] : adj[x]) {
                   ^
transport.cpp: In function 'll setupsuff(int, int, ll, ll)':
transport.cpp:61:1: warning: no return statement in function returning non-void [-Wreturn-type]
 }
 ^
transport.cpp: In function 'll solve()':
transport.cpp:67:28: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for (int i = 0, j = 0; i < suff.size(); ++i) {
                          ~~^~~~~~~~~~~~~
transport.cpp:68:14: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     while (j < pref.size() && suff[i] > pref[j]) {
            ~~^~~~~~~~~~~~~
transport.cpp: In function 'int main()':
transport.cpp:104:21: warning: unused variable 'w' [-Wunused-variable]
     for (auto &[x, w] : adj[c]) {
                     ^

• 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