Submission #1036525

#	Time	Username	Problem	Language	Result	Execution time	Memory
1036525		model_code	Petrol stations (CEOI24_stations)	C++17	100 / 100	1486 ms	1759084 KiB

Main

// 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 (stderr)

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);
      |   ~~~~~^~~~~~~~~~~~~~~~~~~~~~~~

• Subtask 0: 0 / 0
• Subtask 1: 18 / 18
• Subtask 2: 8 / 8
• Subtask 3: 10 / 10
• Subtask 4: 12 / 12
• Subtask 5: 17 / 17
• Subtask 6: 35 / 35