Submission #1036525

# Submission time Handle Problem Language Result Execution time Memory
1036525 2024-07-27T13:31:30 Z model_code Petrol stations (CEOI24_stations) C++17
100 / 100
1486 ms 1759084 KB
// 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