#include <bits/stdc++.h>
using namespace std;
const int kN = 1e5 + 7;
const int kBIT = 2e5 + 7;
int cnt;
class BinaryIndexedTree{
private:
int 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, int delta){
for(; idx <= cnt; 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;
}
};
inline long long min(const long long &a, const long long &b){
return (a < b ? a : b);
}
inline long long max(const long long &a, const long long &b){
return (a > b ? a : b);
}
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 ], cnt);
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);
}
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:82:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:82: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:92:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:92: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:110:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[centroid]){
^
transport.cpp:110: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:129: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:148:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:148: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:162:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:162: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:177:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:188:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:188:22: warning: unused variable 'weight' [-Wunused-variable]
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:195:11: warning: decomposition declaration only available with -std=c++1z or -std=gnu++1z
for(auto [to, weight]: adjacent[node]){
^
transport.cpp:195:22: warning: unused variable 'weight' [-Wunused-variable]
for(auto [to, weight]: adjacent[node]){
^
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
30 ms |
15104 KB |
Output is correct |
| 2 |
Correct |
28 ms |
15488 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
34 ms |
15488 KB |
Output is correct |
| 2 |
Correct |
41 ms |
15920 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
357 ms |
26080 KB |
Output is correct |
| 2 |
Correct |
319 ms |
25128 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
483 ms |
30328 KB |
Output is correct |
| 2 |
Correct |
580 ms |
32356 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
712 ms |
36176 KB |
Output is correct |
| 2 |
Correct |
939 ms |
41720 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
246 ms |
19960 KB |
Output is correct |
| 2 |
Correct |
144 ms |
19868 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
257 ms |
24184 KB |
Output is correct |
| 2 |
Correct |
374 ms |
23544 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
599 ms |
27768 KB |
Output is correct |
| 2 |
Correct |
573 ms |
29944 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
815 ms |
31608 KB |
Output is correct |
| 2 |
Correct |
882 ms |
33560 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Execution timed out |
1054 ms |
35416 KB |
Time limit exceeded |
| 2 |
Halted |
0 ms |
0 KB |
- |
