#include <bits/stdc++.h>
using namespace std;
const int kN = 1e5 + 7;
const int kBIT = 2e5 + 7;
class BinaryIndexedTree{
private:
long long tr[kBIT];
long long query(int idx){
long long ans = 0;
for(; idx >= 1; idx -= idx & -idx)
ans += tr[idx];
return ans;
}
public:
BinaryIndexedTree(){}
void update(int idx, long long delta){
for(; idx < kBIT; idx += idx & -idx)
tr[idx] += delta;
}
long long query(int l, int r){
return query(r) - query(l - 1);
}
};
struct Edge{
int to, weight;
Edge(){}
Edge(int to, int weight){
this->to = to;
this->weight = weight;
}
};
int a[kN], subtree_size[kN], n;
vector<Edge> adjacent[kN];
bool visited[kN];
int centroid_decomposition_root;
pair<long long, long long> prefix[kN], suffix[kN];
map<long long, int> prefix_cnt[kN], suffix_cnt[kN];
vector<int> centroid_decomposition[kN];
set<long long> numbers;
map<long long, int> idx;
BinaryIndexedTree fenwick_prefix, fenwick_suffix;
void read_input(){
ios::sync_with_stdio(false);
cin.tie(NULL);
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;
adjacent[u].push_back(Edge{v, w});
adjacent[v].push_back(Edge{u, w});
}
}
void build_subtree_size(int node = 1, int parent = -1){
subtree_size[node] = 1;
for(auto [to, weight]: adjacent[node]){
if(to == parent || visited[to])
continue;
build_subtree_size(to, node);
subtree_size[node] += subtree_size[to];
}
}
int get_centroid(int node, int parent, const int &component_size){
for(auto [to, weight]: adjacent[node]){
if(!visited[to] && to != parent && subtree_size[to] > component_size / 2){
return get_centroid(to, node, component_size);
}
}
return node;
}
void build_centroid_decomposition(int node = 1, int parent = -1){
build_subtree_size(node);
int centroid = get_centroid(node, -1, subtree_size[node]);
if(parent != -1)
centroid_decomposition[parent].push_back(centroid);
else
centroid_decomposition_root = centroid;
visited[centroid] = true;
for(auto [to, weight]: adjacent[centroid]){
if(!visited[to]){
build_centroid_decomposition(to, centroid);
}
}
visited[centroid] = false;
}
void build_prefix_and_suffix(int node, int weight_to_parent, int parent){
suffix[node].first = suffix[parent].first - weight_to_parent + a[node];
suffix[node].second = min(suffix[parent].first - weight_to_parent, suffix[parent].second);
prefix[node].first = prefix[parent].first - weight_to_parent + a[node];
prefix[node].second = min((long long)a[node] - weight_to_parent, a[node] - weight_to_parent + prefix[parent].second);
prefix[node].second = min(prefix[node].second, 0ll);
array<long long, 2> v{prefix[node].first, -suffix[node].second};
for(auto x: v)
numbers.insert(x);
for(auto [to, weight]: adjacent[node]){
if(to == parent || visited[to])
continue;
build_prefix_and_suffix(to, weight, node);
}
}
long long get_answer_for_subtree(int node, int parent){
long long ans = 0;
ans += prefix[node].second >= 0;
//if prefix can reach to root, adding the a[root] is not necessary
ans += suffix[node].second >= 0;
if(prefix[node].second >= 0)
ans += fenwick_suffix.query(1, idx[ prefix[node].first ]);
ans += fenwick_prefix.query(idx[ -suffix[node].second ], kBIT - 1);
for(auto [to, weight]: adjacent[node]){
if(to == parent || visited[to])
continue;
ans += get_answer_for_subtree(to, node);
}
return ans;
}
void add_to_fenwick(int node, int parent, int delta){
if(prefix[node].second >= 0)
fenwick_prefix.update(idx[ prefix[node].first ], delta);
fenwick_suffix.update(idx[ -suffix[node].second ], delta);
for(auto [to, weight]: adjacent[node]){
if(to == parent || visited[to])
continue;
add_to_fenwick(to, node, delta);
}
}
long long get_answer(int node){
visited[node] = true;
prefix[node] = {0, 0};
suffix[node] = {a[node], 0};
numbers.clear();
idx.clear();
for(auto [to, weight]: adjacent[node]){
if(visited[to])
continue;
build_prefix_and_suffix(to, weight, node);
}
int cnt = 1;
for(long long x: numbers)
idx[x] = cnt++;
long long ans = 0;
for(auto [to, weight]: adjacent[node]){
if(visited[to])
continue;
ans += get_answer_for_subtree(to, node);
add_to_fenwick(to, node, 1);
}
for(auto [to, weight]: adjacent[node]){
if(visited[to])
continue;
add_to_fenwick(to, node, -1);
}
for(int to: centroid_decomposition[node])
ans += get_answer(to);
return ans;
}
int main(){
read_input();
build_centroid_decomposition();
cout << get_answer(centroid_decomposition_root) << "\n";
}
Compilation message
transport.cpp: In function 'void build_subtree_size(int, int)':
transport.cpp:74:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:74:22: warning: unused variable 'weight' [-Wunused-variable]
for(auto [to, weight]: adjacent[node]){
^
transport.cpp: In function 'int get_centroid(int, int, const int&)':
transport.cpp:84:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:84:22: warning: unused variable 'weight' [-Wunused-variable]
for(auto [to, weight]: adjacent[node]){
^
transport.cpp: In function 'void build_centroid_decomposition(int, int)':
transport.cpp:102:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[centroid]){
^
transport.cpp:102:22: warning: unused variable 'weight' [-Wunused-variable]
for(auto [to, weight]: adjacent[centroid]){
^
transport.cpp: In function 'void build_prefix_and_suffix(int, int, int)':
transport.cpp:121:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp: In function 'long long int get_answer_for_subtree(int, int)':
transport.cpp:140:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:140:22: warning: unused variable 'weight' [-Wunused-variable]
for(auto [to, weight]: adjacent[node]){
^
transport.cpp: In function 'void add_to_fenwick(int, int, int)':
transport.cpp:154:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:154:22: warning: unused variable 'weight' [-Wunused-variable]
for(auto [to, weight]: adjacent[node]){
^
transport.cpp: In function 'long long int get_answer(int)':
transport.cpp:169:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:180:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:180:22: warning: unused variable 'weight' [-Wunused-variable]
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:187:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:187:22: warning: unused variable 'weight' [-Wunused-variable]
for(auto [to, weight]: adjacent[node]){
^
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
29 ms |
15232 KB |
Output is correct |
| 2 |
Correct |
28 ms |
15616 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
38 ms |
15616 KB |
Output is correct |
| 2 |
Correct |
44 ms |
16028 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
389 ms |
26232 KB |
Output is correct |
| 2 |
Correct |
349 ms |
25208 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
514 ms |
30712 KB |
Output is correct |
| 2 |
Correct |
640 ms |
34172 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
767 ms |
36400 KB |
Output is correct |
| 2 |
Correct |
994 ms |
44920 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
280 ms |
20216 KB |
Output is correct |
| 2 |
Correct |
146 ms |
20096 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
277 ms |
24568 KB |
Output is correct |
| 2 |
Correct |
391 ms |
23672 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
609 ms |
28288 KB |
Output is correct |
| 2 |
Correct |
647 ms |
32376 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
858 ms |
32120 KB |
Output is correct |
| 2 |
Correct |
884 ms |
36344 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1098 ms |
35704 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
