oj.uz

Submission #1049162

# Submission time Handle Problem Language Result Execution time Memory
1049162 2024-08-08T14:23:19 Z NeroZein Transport (COCI19_transport) C++17
130 / 130
 422 ms 25680 KB 
#include "bits/stdc++.h"
#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>

using namespace std;
using namespace __gnu_pbds;

#ifdef Nero
#include "Deb.h"
#else
#define debug(...)
#endif

const int N = 1e5 + 5;

template <class T> using Tree = tree<T,null_type, less_equal<T>, rb_tree_tag,tree_order_statistics_node_update>;

int a[N];
int sz[N];
bool block[N];
long long sum[N];
vector<pair<int, int>> g[N];
long long min_up[N], min_down[N];

void find_sizes(int v, int p) {
  sz[v] = 1;
  for (auto [u, w] : g[v]) {
    if (!block[u] && u != p) {
      find_sizes(u, v);
      sz[v] += sz[u];
    }
  }
}

int find_centroid(int v, int p, int x) {
  for (auto [u, w] : g[v]) {
    if (!block[u] && u != p && sz[u] > x / 2) {
      return find_centroid(u, v, x); 
    }
  }
  return v;
}

long long dfs_qry(int v, int p, Tree<long long>& se, const int centroid, const bool rep) {
  long long ret = 0; 
  if (min_down[v] >= 0 && !rep) {
    ret++;//centroid -> v
  }
  if (min_up[v] >= 0) {
    ret += (int) se.size() - se.order_of_key(-(sum[v] - a[centroid]));    
  }
  for (auto [u, w] : g[v]) {
    if (u == p || block[u]) {
      continue; 
    }
    sum[u] = sum[v] + a[u] - w; 
    min_up[u] = min(0LL, min_up[v] + a[u] - w);
    min_down[u] = min(min_down[v], sum[v] - w);
    ret += dfs_qry(u, v, se, centroid, rep); 
  }
  return ret; 
}

void dfs_upd(int v, int p, Tree<long long>& se) {
  se.insert(min_down[v]);
  for (auto [u, w] : g[v]) {
    if (u != p && !block[u]) {
      dfs_upd(u, v, se); 
    }
  }
}

long long decompose(int node) {
  
  find_sizes(node, node);
  int centroid = find_centroid(node, node, sz[node]);
  block[centroid] = true; 
  
  long long ret = 0;
  for (int rep = 0; rep < 2; ++rep) {
    Tree<long long> se; 
    if (rep == 0) {
      se.insert(0);      
    }
    min_down[centroid] = min_up[centroid] = 0;
    sum[centroid] = a[centroid];
    for (auto [u, w] : g[centroid]) {
      if (block[u]) {
        continue; 
      }
      sum[u] = sum[centroid] + a[u] - w;
      min_up[u] = min(0, a[u] - w); 
      min_down[u] = min(0, a[centroid] - w);
      ret += dfs_qry(u, centroid, se, centroid, rep);
      dfs_upd(u, centroid, se);
    }
    reverse(g[centroid].begin(), g[centroid].end());    
  }
  
  for (auto [u, w] : g[centroid]) {
    if (!block[u]) {
      ret += decompose(u); 
    }
  }
  return ret;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  //freopen("transport4.in", "r", stdin); 
  int n;
  cin >> n;
  for (int i = 1; i <= n; ++i) {
    cin >> a[i];
  }
  for(int i = 0; i < n - 1; ++i) {
    int u, v, w;
    cin >> u >> v >> w;
    g[u].push_back({v, w});
    g[v].push_back({u, w});
  }
  long long ans = decompose(1);
  cout << ans << '\n';
  return 0;
}

Subtask #1
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 8 ms 4928 KB Output is correct
2 Correct 6 ms 4956 KB Output is correct

Subtask #2
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 10 ms 5336 KB Output is correct
2 Correct 13 ms 5292 KB Output is correct

Subtask #3
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 161 ms 12464 KB Output is correct
2 Correct 148 ms 12744 KB Output is correct

Subtask #4
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 211 ms 15440 KB Output is correct
2 Correct 240 ms 18688 KB Output is correct

Subtask #5
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 295 ms 19536 KB Output is correct
2 Correct 416 ms 25680 KB Output is correct

Subtask #6
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 91 ms 7660 KB Output is correct
2 Correct 44 ms 7256 KB Output is correct

Subtask #7
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 76 ms 9816 KB Output is correct
2 Correct 143 ms 9172 KB Output is correct

Subtask #8
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 183 ms 11344 KB Output is correct
2 Correct 177 ms 11600 KB Output is correct

Subtask #9
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 248 ms 13188 KB Output is correct
2 Correct 268 ms 13652 KB Output is correct

Subtask #10
13.0 / 13.0

# Verdict Execution time Memory Grader output
1 Correct 422 ms 16212 KB Output is correct
2 Correct 390 ms 15240 KB Output is correct