// Author: Jiří Kalvoda
#include<bits/stdc++.h>
using namespace std;
using ll = long long;
int n;
ll k;
struct Treap
{
Treap *left=0, *right=0;
ll key;
int priority;
int count;
int subtree_count;
int subtree_size;
Treap * clone();
Treap * calc()
{
subtree_count = count;
if (left) subtree_count += left->subtree_count;
if (right) subtree_count += right->subtree_count;
subtree_size = 1;
if (left) subtree_size += left->subtree_size;
if (right) subtree_size += right->subtree_size;
return this;
}
Treap * init(ll key_, int count_)
{
this->priority = random();
this->key = key_;
this->count = count_;
calc();
return this;
}
};
Treap * treap_aloc;
int treap_aloc_count;
Treap *new_treap()
{
if(!treap_aloc_count)
{
treap_aloc_count = 1000;
treap_aloc = new Treap[treap_aloc_count];
}
treap_aloc_count--;
return treap_aloc++;
}
Treap * Treap::clone()
{
Treap * t = new_treap();
t->left = left;
t->right = right;
t->key = key;
t->count = count;
t->priority = priority;
t->calc();
return t;
}
int max_h = 0;
Treap *merge(Treap *left, Treap *right, int h=0)
{
max_h = max(max_h, h);
if (!left || !right) return left?left:right;
if (left->priority < right->priority)
{
left = left->clone();
left->right = merge(left->right, right, h+1);
return left->calc();
}
else
{
right = right->clone();
right->left = merge(left, right->left, h+1);
return right->calc();
}
}
pair<Treap*, Treap*> split(Treap *in, ll key, int h=0)
{
max_h = max(max_h, h);
if (!in) return {NULL, NULL};
in = in->clone();
Treap *left=in,*right=in;
if (in->key < key)
tie(left->right, right) = split(in->right, key, h+1);
else
tie(left, right->left) = split(in->left, key, h+1);
if(left) left = left->calc();
if(right) right = right->calc();
return {left, right};
}
void printttr(FILE * out, Treap *tr, ll key_shift)
{
if(!tr) return;
fprintf(out, "{");
printttr(out, tr->left, key_shift);
fprintf(out, "[%d %d %d]", tr->key+key_shift, tr->count, tr->subtree_count);
printttr(out, tr->right, key_shift);
fprintf(out, "}");
}
struct DataStr {
Treap * tr=NULL;
ll key_shift=0;
int del_negative()
{
Treap *to_remove;
tie(to_remove, tr) = split(tr, 0-key_shift);
return to_remove?to_remove->subtree_count:0;
}
void add(ll key, int count)
{
if(!count) return;
Treap *a, *b, *c;
tie(a,b) = split(tr, key-key_shift);
tie(b,c) = split(b, key-key_shift+1);
if(b)
{
b = b->clone();
b->count += count;
b->calc();
}
else
b = new_treap()->init(key-key_shift, count);
tr = merge(a, merge(b,c));
}
void add_from(Treap * from, ll from_key_shift, int multiple=1)
{
if(!from) return;
add(from->key + from_key_shift, from->count * multiple);
add_from(from->left, from_key_shift, multiple);
add_from(from->right, from_key_shift, multiple);
}
};
DataStr merge_ds(DataStr a, DataStr b)
{
if(!a.tr) return b;
if(!b.tr) return a;
if(a.tr->subtree_size > b.tr->subtree_size)
return merge_ds(b, a);
b.add_from(a.tr, a.key_shift);
return b;
}
struct Node{
vector<pair<Node*, ll>> e;
pair<Node*, ll> up;
ll counts=0;
int id;
int subtree_size;
DataStr dp_here;
DataStr dp_after_edge;
int make_rooted(Node * _up)
{
subtree_size=1;
for(int i=0; i<e.size(); i++)
{
if(e[i].first == _up)
{
up = e[i];
e[i] = e[e.size()-1];
e.pop_back();
}
}
for(auto &[nd, l] : e)
subtree_size += nd->make_rooted(this);
sort(e.begin(), e.end(), [](auto &a, auto &b){return a.first->subtree_size > b.first->subtree_size;});
return subtree_size;
}
void go()
{
dp_here.add(k, 1);
for(auto &[nd, l] : e)
nd->go();
for(auto &[nd, l] : e)
dp_here = merge_ds(nd->dp_after_edge, dp_here);
dp_after_edge = dp_here;
dp_after_edge.key_shift -= up.second;
ll c = dp_after_edge.del_negative();
counts += c*(n-subtree_size);
dp_after_edge.add(k-up.second, c);
// printf("GO %d %d\n", id, c);
// printttr(stdout, dp_after_edge.tr, dp_after_edge.key_shift);
// printf("\n");
}
void go2(vector<DataStr> dp_down)
{
bool first = true;
for(auto &[nd, l] : e)
{
vector<DataStr> dp_down_nd = dp_down;
assert(1<<dp_down_nd.size() <= n);
if(first)
{
first=false;
dp_down_nd[0].add(k, 1);
for(auto &[nd2, l2] : e) if(nd2 != nd)
dp_down_nd[0] = merge_ds(dp_down_nd[0], nd2->dp_after_edge);
}
else
{
DataStr tmp = dp_here;
tmp.add_from(nd->dp_after_edge.tr, nd->dp_after_edge.key_shift, -1);
dp_down_nd.push_back(tmp);
}
ll c = 0;
for(auto &it: dp_down_nd)
{
it.key_shift -= l;
c += it.del_negative();
}
counts += c*nd->subtree_size;
dp_down_nd[0].add(k-l, c);
nd->go2(dp_down_nd);
}
}
} *nds;
int main(int argc, char ** argv)
{
srand(42);
scanf("%d%lld", &n, &k);
nds = new Node[n];
for (int i=0; i<n; i++) nds[i].id = i;
for (int i=0; i<n-1; i++)
{
int a,b;
ll l;
scanf("%d%d%lld", &a, &b, &l);
nds[a].e.push_back({nds+b, l});
nds[b].e.push_back({nds+a, l});
}
nds[0].make_rooted(NULL);
nds[0].go();
nds[0].go2({DataStr()});
for (int i=0; i<n; i++) printf("%lld\n", nds[i].counts);
fprintf(stderr, "max treap height: %d\n", max_h);
return 0;
}
Compilation message
Main.cpp: In function 'void printttr(FILE*, Treap*, ll)':
Main.cpp:103:18: warning: format '%d' expects argument of type 'int', but argument 3 has type 'll' {aka 'long long int'} [-Wformat=]
103 | fprintf(out, "[%d %d %d]", tr->key+key_shift, tr->count, tr->subtree_count);
| ~^ ~~~~~~~~~~~~~~~~~
| | |
| int ll {aka long long int}
| %lld
Main.cpp: In member function 'int Node::make_rooted(Node*)':
Main.cpp:163:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<Node*, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
163 | for(int i=0; i<e.size(); i++)
| ~^~~~~~~~~
Main.cpp: In function 'int main(int, char**)':
Main.cpp:233:7: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
233 | scanf("%d%lld", &n, &k);
| ~~~~~^~~~~~~~~~~~~~~~~~
Main.cpp:240:8: warning: ignoring return value of 'int scanf(const char*, ...)' declared with attribute 'warn_unused_result' [-Wunused-result]
240 | scanf("%d%d%lld", &a, &b, &l);
| ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
344 KB |
Output is correct |
2 |
Correct |
1 ms |
348 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
344 KB |
Output is correct |
2 |
Correct |
1 ms |
348 KB |
Output is correct |
3 |
Correct |
4 ms |
5720 KB |
Output is correct |
4 |
Correct |
2 ms |
4188 KB |
Output is correct |
5 |
Correct |
3 ms |
3288 KB |
Output is correct |
6 |
Correct |
1 ms |
1628 KB |
Output is correct |
7 |
Correct |
1 ms |
1112 KB |
Output is correct |
8 |
Correct |
1 ms |
344 KB |
Output is correct |
9 |
Correct |
2 ms |
2140 KB |
Output is correct |
10 |
Correct |
2 ms |
2136 KB |
Output is correct |
11 |
Correct |
4 ms |
2140 KB |
Output is correct |
12 |
Correct |
4 ms |
2140 KB |
Output is correct |
13 |
Correct |
3 ms |
2396 KB |
Output is correct |
14 |
Correct |
6 ms |
4700 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
344 KB |
Output is correct |
2 |
Correct |
173 ms |
72484 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
344 KB |
Output is correct |
2 |
Correct |
1 ms |
348 KB |
Output is correct |
3 |
Correct |
1 ms |
344 KB |
Output is correct |
4 |
Correct |
173 ms |
72484 KB |
Output is correct |
5 |
Correct |
139 ms |
87668 KB |
Output is correct |
6 |
Correct |
328 ms |
352848 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
344 KB |
Output is correct |
2 |
Correct |
1 ms |
348 KB |
Output is correct |
3 |
Correct |
141 ms |
87384 KB |
Output is correct |
4 |
Correct |
121 ms |
83688 KB |
Output is correct |
5 |
Correct |
86 ms |
51388 KB |
Output is correct |
6 |
Correct |
0 ms |
348 KB |
Output is correct |
7 |
Correct |
155 ms |
87008 KB |
Output is correct |
8 |
Correct |
154 ms |
86732 KB |
Output is correct |
9 |
Correct |
146 ms |
86868 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
344 KB |
Output is correct |
2 |
Correct |
1 ms |
348 KB |
Output is correct |
3 |
Correct |
141 ms |
87384 KB |
Output is correct |
4 |
Correct |
121 ms |
83688 KB |
Output is correct |
5 |
Correct |
86 ms |
51388 KB |
Output is correct |
6 |
Correct |
0 ms |
348 KB |
Output is correct |
7 |
Correct |
155 ms |
87008 KB |
Output is correct |
8 |
Correct |
154 ms |
86732 KB |
Output is correct |
9 |
Correct |
146 ms |
86868 KB |
Output is correct |
10 |
Correct |
142 ms |
106436 KB |
Output is correct |
11 |
Correct |
137 ms |
102696 KB |
Output is correct |
12 |
Correct |
127 ms |
92788 KB |
Output is correct |
13 |
Correct |
164 ms |
108112 KB |
Output is correct |
14 |
Correct |
143 ms |
107600 KB |
Output is correct |
15 |
Correct |
178 ms |
108200 KB |
Output is correct |
16 |
Correct |
139 ms |
106688 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
344 KB |
Output is correct |
2 |
Correct |
1 ms |
348 KB |
Output is correct |
3 |
Correct |
4 ms |
5720 KB |
Output is correct |
4 |
Correct |
2 ms |
4188 KB |
Output is correct |
5 |
Correct |
3 ms |
3288 KB |
Output is correct |
6 |
Correct |
1 ms |
1628 KB |
Output is correct |
7 |
Correct |
1 ms |
1112 KB |
Output is correct |
8 |
Correct |
1 ms |
344 KB |
Output is correct |
9 |
Correct |
2 ms |
2140 KB |
Output is correct |
10 |
Correct |
2 ms |
2136 KB |
Output is correct |
11 |
Correct |
4 ms |
2140 KB |
Output is correct |
12 |
Correct |
4 ms |
2140 KB |
Output is correct |
13 |
Correct |
3 ms |
2396 KB |
Output is correct |
14 |
Correct |
6 ms |
4700 KB |
Output is correct |
15 |
Correct |
1 ms |
344 KB |
Output is correct |
16 |
Correct |
173 ms |
72484 KB |
Output is correct |
17 |
Correct |
139 ms |
87668 KB |
Output is correct |
18 |
Correct |
328 ms |
352848 KB |
Output is correct |
19 |
Correct |
141 ms |
87384 KB |
Output is correct |
20 |
Correct |
121 ms |
83688 KB |
Output is correct |
21 |
Correct |
86 ms |
51388 KB |
Output is correct |
22 |
Correct |
0 ms |
348 KB |
Output is correct |
23 |
Correct |
155 ms |
87008 KB |
Output is correct |
24 |
Correct |
154 ms |
86732 KB |
Output is correct |
25 |
Correct |
146 ms |
86868 KB |
Output is correct |
26 |
Correct |
142 ms |
106436 KB |
Output is correct |
27 |
Correct |
137 ms |
102696 KB |
Output is correct |
28 |
Correct |
127 ms |
92788 KB |
Output is correct |
29 |
Correct |
164 ms |
108112 KB |
Output is correct |
30 |
Correct |
143 ms |
107600 KB |
Output is correct |
31 |
Correct |
178 ms |
108200 KB |
Output is correct |
32 |
Correct |
139 ms |
106688 KB |
Output is correct |
33 |
Correct |
1486 ms |
1759084 KB |
Output is correct |
34 |
Correct |
338 ms |
322760 KB |
Output is correct |
35 |
Correct |
208 ms |
130052 KB |
Output is correct |
36 |
Correct |
195 ms |
126844 KB |
Output is correct |
37 |
Correct |
975 ms |
1292072 KB |
Output is correct |
38 |
Correct |
1058 ms |
1288884 KB |
Output is correct |
39 |
Correct |
1108 ms |
1335932 KB |
Output is correct |
40 |
Correct |
516 ms |
551876 KB |
Output is correct |