#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
template<class K, class V>
using ordered_map = __gnu_pbds::tree<K, V, std::less<K>, __gnu_pbds::rb_tree_tag, __gnu_pbds::tree_order_statistics_node_update>;
template<class K>
using ordered_set = ordered_map<K, __gnu_pbds::null_type>;
namespace std {
template<class Fun>
class y_combinator_result {
Fun fun_;
public:
template<class T>
explicit y_combinator_result(T &&fun): fun_(std::forward<T>(fun)) {}
template<class ...Args>
decltype(auto) operator()(Args &&...args) {
return fun_(std::ref(*this), std::forward<Args>(args)...);
}
};
template<class Fun>
decltype(auto) y_combinator(Fun &&fun) {
return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun));
}
} // namespace std
int main() {
using namespace std;
ios_base::sync_with_stdio(false);
cin.tie(NULL);
int N;
cin >> N;
vector<long long> A(N);
for (auto &a : A)
cin >> a;
vector<vector<pair<int, long long>>> T(N);
for (int i = 1; i < N; i++) {
int u, v;
long long w;
cin >> u >> v >> w;
u--, v--;
T[u].emplace_back(v, w);
T[v].emplace_back(u, w);
}
vector<bool> blk(N, false);
vector<int> wei(N);
auto init_wei = y_combinator([&](auto self, int u, int p = -1) -> int {
wei[u] = 1;
for (auto [v, w] : T[u]) {
if (v == p || blk[v])
continue;
wei[u] += self(v, u);
}
return wei[u];
});
auto find_cen = y_combinator([&](auto self, int n, int u, int p = -1) -> int {
for (auto [v, w] : T[u]) {
if (v == p || blk[v])
continue;
if (2 * wei[v] > n)
return self(n, v, u);
}
return u;
});
// compute fuel from u to root
auto dfs0 = y_combinator([&](auto self, long long tot, long long mgp, vector<long long> &src, int u, int p) -> void {
if (mgp >= 0)
src.push_back(tot);
for (auto [v, w] : T[u]) {
if (v == p || blk[v])
continue;
auto ntot = tot - w + A[v];
auto nmgp = min(mgp, 0LL) - w + A[v];
self(ntot, nmgp, src, v, u);
}
});
// compute fuel from root to u
auto dfs1 = y_combinator([&](auto self, long long cur, long long ret, vector<long long> &snk, int u, int p) -> void {
snk.push_back(ret);
cur += A[u];
for (auto [v, w] : T[u]) {
if (v == p || blk[v])
continue;
if (cur < w)
self(0, ret + w - cur, snk, v, u);
else
self(cur - w, ret, snk, v, u);
}
});
cout << y_combinator([&](auto self, int u = 0) -> long long {
int n = init_wei(u);
u = find_cen(n, u);
int ind = 0;
long long ret = 0;
ordered_set<pair<long long, int>> ord1, ord2;
ord1.insert(make_pair(0, ind++));
ord2.insert(make_pair(0, ind++));
for (auto [v, w] : T[u]) {
if (blk[v])
continue;
vector<long long> src, snk;
dfs0(A[v] - w, A[v] - w, src, v, u);
dfs1(max(A[u] - w, 0LL), max(w - A[u], 0LL), snk, v, u);
long long tem = 0;
for (auto x : src)
tem += ord2.order_of_key(make_pair(x, INT_MAX));
for (auto x : snk)
tem += (int)ord1.size() - ord1.order_of_key(make_pair(x, INT_MIN));
for (auto x : src)
ord1.insert(make_pair(x, ind++));
for (auto x : snk)
ord2.insert(make_pair(x, ind++));
ret += tem;
}
blk[u] = true;
for (auto [v, w] : T[u])
if (!blk[v])
ret += self(v);
return ret;
})() << '\n';
}
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
11 ms |
980 KB |
Output is correct |
2 |
Correct |
8 ms |
1236 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
16 ms |
1608 KB |
Output is correct |
2 |
Correct |
17 ms |
1620 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
238 ms |
14544 KB |
Output is correct |
2 |
Correct |
172 ms |
11460 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
288 ms |
18468 KB |
Output is correct |
2 |
Correct |
326 ms |
19092 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
435 ms |
25628 KB |
Output is correct |
2 |
Correct |
480 ms |
28248 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
160 ms |
8520 KB |
Output is correct |
2 |
Correct |
64 ms |
6152 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
131 ms |
11640 KB |
Output is correct |
2 |
Correct |
200 ms |
10380 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
262 ms |
14704 KB |
Output is correct |
2 |
Correct |
242 ms |
14400 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
357 ms |
20268 KB |
Output is correct |
2 |
Correct |
341 ms |
18732 KB |
Output is correct |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
640 ms |
26716 KB |
Output is correct |
2 |
Correct |
509 ms |
22264 KB |
Output is correct |