#include <bits/stdc++.h>
#define all(x) (x).begin(), (x).end()
#define ff first
#define ss second
#define pb push_back
#define mp make_pair
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef long double ld;
typedef pair<ll, ll> pll;
typedef pair<ull, ull> pull;
typedef pair<int, int> pii;
typedef pair<ld, ld> pld;
int n, k;
struct fenwick{
vector<int> bit;
vector<int> vals;
void init(){
sort(all(vals));
vals.resize(unique(all(vals))-vals.begin());
bit.resize(vals.size()+1);
}
void upd(int idx, int val){
idx = lower_bound(all(vals), idx)-vals.begin();
for(++idx; idx < bit.size(); idx += idx&(-idx))
bit[idx] += val;
}
int get(int idx, int smaller = 0){
idx = lower_bound(all(vals), idx)-vals.begin();
idx -= smaller;
if(idx == vals.size()) --idx;
int ret = 0;
for(++idx; idx > 0; idx -= idx&(-idx))
ret += bit[idx];
return ret;
}
};
vector<pii> g[100009];
int subsz[100009];
int block[100009];
void getsz(int v, int prt){
subsz[v] = 1;
for(auto u : g[v])
if(u.ff != prt && !block[u.ff]){
getsz(u.ff, v);
subsz[v] += subsz[u.ff];
}
}
int getcentroid(int v, int prt, int tot){
for(auto u : g[v])
if(u.ff != prt && !block[u.ff] && subsz[u.ff]*2 >= tot)
return getcentroid(u.ff, v, tot);
return v;
}
vector<pll> child[100009];
void dfs(int v, int prt,int d, int curmax, int top){
child[top].pb({curmax, d});
for(auto u : g[v])
if(u.ff != prt && !block[u.ff])
dfs(u.ff, v, d+1, max(curmax, u.ss), top);
}
vector<int> rel;
fenwick seg[400009];
void add(int v, int tl, int tr, int idx, int val){
rel.pb(v);
seg[v].vals.pb(val);
if(tl == tr) return;
if(idx <= (tl+tr)/2)
add(2*v, tl, (tl+tr)/2, idx, val);
else
add(2*v+1, (tl+tr)/2+1, tr, idx, val);
}
void upd(int v, int tl, int tr, int idx, int val, int change){
seg[v].upd(val, change);
if(tl == tr) return;
if(idx <= (tl+tr)/2)
upd(2*v, tl, (tl+tr)/2, idx, val, change);
else
upd(2*v+1, (tl+tr)/2+1, tr, idx, val, change);
}
ll mx = 1e5;
ll ans;
ll doubleans;
void get(int v, int tl, int tr, int l, int r, int val){
if(tr < l || r < tl)
return;
if(l <= tl && tr <= r){
assert(seg[v].vals.empty() || seg[v].bit.size() == seg[v].vals.size()+1);
int small = seg[v].get(val, 1);
/*if(small < 0){
cout << tl << ' ' << tr << ' ' << val << '\n';
for(auto u : seg[v].vals)
cout << u << ' ';
cout << '\n';
}*/
ans += small;
doubleans += seg[v].get(val)-small;
}
else{
get(2*v, tl, (tl+tr)/2, l, r, val);
get(2*v+1, (tl+tr)/2+1, tr, l, r, val);
}
}
void getans(int v){
getsz(v, 0);
v = getcentroid(v, 0, subsz[v]);
for(auto u : g[v]){
if(block[u.ff]) continue;
dfs(u.ff, v, 1, u.ss, u.ff);
for(auto u : child[u.ff]){
if(u.ff-u.ss-k >= 0){
++ans;
}
add(1, 0, mx, u.ss, u.ff);
}
}
for(auto u : rel)
if(seg[u].bit.empty())
seg[u].init();
for(auto u : g[v]){
if(block[u.ff]) continue;
for(auto u : child[u.ff]){
upd(1, 0, mx, u.ss, u.ff, 1);
}
}
for(auto u : g[v]){
if(block[u.ff]) continue;
for(auto u : child[u.ff]){
upd(1, 0, mx, u.ss, u.ff, -1);
}
for(auto u : child[u.ff]){
get(1, 0, mx, 0, u.ff-u.ss-k, u.ff);
}
for(auto u : child[u.ff]){
upd(1, 0, mx, u.ss, u.ff, 1);
}
}
for(auto u : rel){
seg[u].bit.clear();
seg[u].vals.clear();
}
for(auto u : g[v])
child[u.ff].clear();
rel.clear();
block[v] = 1;
for(auto u : g[v])
if(!block[u.ff])
getans(u.ff);
}
int main() {
ios_base::sync_with_stdio(false);
cin.tie(nullptr);
//freopen("in.txt", "r", stdin);
//freopen("out.txt", "w", stdout);
cin >> n >> k;
for(int i = 0; i < n-1; ++i){
int x, y, z;
cin >> x >> y >> z;
g[x].pb({y, z});
g[y].pb({x, z});
}
getans(1);
cout << ans*2+doubleans << '\n';
}
Compilation message
Main.cpp: In member function 'void fenwick::upd(int, int)':
Main.cpp:26:24: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
26 | for(++idx; idx < bit.size(); idx += idx&(-idx))
| ~~~~^~~~~~~~~~~~
Main.cpp: In member function 'int fenwick::get(int, int)':
Main.cpp:32:16: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
32 | if(idx == vals.size()) --idx;
| ~~~~^~~~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
18 ms |
23788 KB |
Output is correct |
| 2 |
Correct |
16 ms |
23788 KB |
Output is correct |
| 3 |
Incorrect |
18 ms |
23916 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
17 ms |
23788 KB |
Output is correct |
| 2 |
Correct |
16 ms |
23788 KB |
Output is correct |
| 3 |
Incorrect |
17 ms |
23916 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
18 ms |
23788 KB |
Output is correct |
| 2 |
Correct |
16 ms |
23788 KB |
Output is correct |
| 3 |
Incorrect |
18 ms |
23916 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
