Submission #1047842

#	Time	Username	Problem	Language	Result	Execution time	Memory
1047842		0npata	Beech Tree (IOI23_beechtree)	C++17	100 / 100	131 ms	91352 KiB

beechtree

#include "beechtree.h"
#include<bits/stdc++.h>
using namespace std;

#define vec vector

const int MXN = 200'005;

vec<int> tree[MXN];

map<int, int> tree_col_sz[MXN];
map<int, int> tree_ord[MXN];

vec<int> ans;
int tree_sz[MXN];

void dfs0(int u = 0) {
	tree_sz[u] = 1;
	for(int v : tree[u]) {
		dfs0(v);
		tree_sz[u] += tree_sz[v];
	}
}

void dfs(int u = 0) {
	tree_ord[u].insert({tree_sz[u], u});
	if(tree[u].size() == 0) {
		return;
	}
	for(int v : tree[u]) {
		dfs(v);
		ans[u] &= ans[v];
	}
	if(ans[u] == 0) return;
	sort(tree[u].begin(), tree[u].end(), [&](int v, int w) { return tree_ord[v].size() > tree_ord[w].size(); });
	tree[u].push_back(u);
	int f = tree[u][0];

	for(int i = 1; i<tree[u].size(); i++) {
		int v = tree[u][i];
		for(auto [w_sz, w] : tree_ord[v]) {
			auto it = tree_ord[f].lower_bound(w_sz);
			if(it != tree_ord[f].end()) {
				int x = it->second;
				for(auto [c, c_cnt] : tree_col_sz[w]) {
					if(tree_col_sz[x][c] < c_cnt) {
						ans[u] = 0;
						return;
					}
				}
			}
			if(it != tree_ord[f].begin()) {
				it--;
				int x = it->second;
				for(auto [c, c_cnt] : tree_col_sz[x]) {
					if(tree_col_sz[w][c] < c_cnt) {
						ans[u] = 0;
						return;
					}
				}
			}
			tree_ord[f].insert({w_sz, w});
		}
	}
	tree_ord[u] = move(tree_ord[f]);
}

std::vector<int> beechtree(int N, int M, std::vector<int> P, std::vector<int> C)
{
	ans = vec<int>(N, 1);

	for(int i = 1; i<N; i++) {
		tree[P[i]].push_back(i);
	}

	dfs0();

	//for(int i = 0; i<N;  i++) cerr << tree_sz[i] << ' ';
	//cerr << '\n';

	for(int i = 1; i<N; i++) {
		if(tree_col_sz[P[i]].count(C[i])) {
			ans[P[i]] = 0;
		}
		tree_col_sz[P[i]][C[i]] = tree_sz[i];
	}

	dfs();

	return ans;
}

Compilation message (stderr)

beechtree.cpp: In function 'void dfs(int)':
beechtree.cpp:39:18: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   39 |  for(int i = 1; i<tree[u].size(); i++) {
      |                 ~^~~~~~~~~~~~~~~

• Subtask 0: 0 / 0
• Subtask 1: 9 / 9
• Subtask 2: 5 / 5
• Subtask 3: 9 / 9
• Subtask 4: 8 / 8
• Subtask 5: 14 / 14
• Subtask 6: 14 / 14
• Subtask 7: 12 / 12
• Subtask 8: 17 / 17
• Subtask 9: 12 / 12