Submission #785941

#TimeUsernameProblemLanguageResultExecution timeMemory
785941KiaRezJanjetina (COCI21_janjetina)C++17
50 / 110
332 ms24432 KiB
/* IN THE NAME OF GOD */ #include <bits/stdc++.h> // #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt") using namespace std; typedef long long ll; typedef pair<ll, ll> pll; typedef pair<int, int> pii; typedef pair<int, pii> ipii; typedef pair<pii, int> piii; typedef pair<ll, pll> lpll; typedef pair<pll, ll> plll; typedef pair<pii, pii> piipii; typedef pair<pll, pll> pllpll; typedef long double ld; #define SQ 350 #define endl '\n' #define F first #define S second #define Mp make_pair #define pb push_back #define pf push_front #define PQ priority_queue #define size(x) ((ll)x.size()) #define all(x) (x).begin(),(x).end() #define kill(x) cout << x << '\n', exit(0); #define set_dec(x) cout << fixed << setprecision(x); #define ios ios_base::sync_with_stdio(false), cin.tie(0) #define fuck(x) cout << "(" << #x << " , " << x << ")" << endl #define file_io(x,y) freopen(x, "r", stdin); freopen(y, "w", stdout); const int N = 4e5+25, lg=18; ll Mod = 1e9+7; //ll Mod = 998244353; ll MOD(ll a, ll mod=Mod) {return (a>=0&&a<mod ? a : ((a%mod+mod)%mod));} ll max(ll a, ll b) {return (a>b ? a : b);} ll min(ll a, ll b) {return (a<b ? a : b);} ll abso(ll a) {return (a<0?-a:a);} inline ll poww(ll a, ll b) { ll ans = 1; a=MOD(a); while (b) { if (b & 1) ans = MOD(ans*a); b >>= 1; a = MOD(a*a); } return ans; } ll ans = 0; int n, k, subt[N], mark[N]; ll c[N]; vector<pii> adj[N]; vector<vector<pii>> vec; inline ll lowbit(int x) { return x & (-x); } inline void update(int x, int d) { while (x <= n+2) { c[x] += d; x += lowbit(x); } } inline ll query(int x) {int res = 0; while(x) { res += c[x]; x -= lowbit(x); } return res; } int findCentroid(int v, int p, int sz) { for(auto u : adj[v]) { if(u.F == p || mark[u.F]) continue; if(subt[u.F] > sz/2) return findCentroid(u.F, v, sz); } return v; } void calcSubt(int v, int p) { subt[v] = 1; for(auto u : adj[v]) { if(mark[u.F] || u.F == p) continue; calcSubt(u.F, v); subt[v] += subt[u.F]; } } void calcCandidates(int v, int p, int mx, int l) { vec.back().pb({mx, l}); if(mx - l >= k) {ans++;} for(auto u : adj[v]) { if(u.F == p || mark[u.F]) continue; calcCandidates(u.F, v, max(u.S, mx), l+1); } } void solve(int v) { calcSubt(v, 0); int cent = findCentroid(v, 0, subt[v]); while(size(vec)) vec.pop_back(); vector<pii> vec2; for(auto u : adj[cent]) { if(mark[u.F]) continue; vec.pb({}); calcCandidates(u.F, cent, u.S, 1); for(auto it : vec.back()) vec2.pb(it); } sort(all(vec2)); for(auto it : vec2) { if(it.F-k-it.S>=0) ans += ((ll)query(it.F-it.S-k)); update(it.S, 1); } for(auto it : vec2) update(it.S, -1); for(int i=0; i<size(vec); i++) { vec2 = vec[i]; sort(all(vec2)); for(auto it : vec2) { if(it.F-it.S-k>=0) ans -= ((ll)query(it.F-it.S-k)); update(it.S, 1); } for(auto it : vec2) update(it.S, -1); } mark[cent] = 1; for(auto u : adj[cent]) { if(mark[u.F]) continue; solve(u.F); } } int main () { ios; cin>>n>>k; for(int v,u,w,i=1; i<n; i++) { cin>>v>>u>>w; adj[v].pb({u, w}); adj[u].pb({v, w}); } solve(1); cout<<2*ans<<endl; return 0; }
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...