Zenn
👏

ABC239 E - Subtree K-th Max

2023/01/18に公開
AtCoder
tech

問題概要

問題文

N頂点の根付き木が与えられる。根は1である。また、頂点iには整数X_iが書かれている。
Q個のクエリが与えられる。i番目のクエリでは、整数の組(V_i,K_i)が与えられるので、次の問題に答えよ。

  • 問題:頂点V_iの部分木に含まれる頂点に書かれた整数のうち、K_i番目に大きいものを求めよ。

制約

  • 1\leq N\leq 10^5
  • 0\leq X\leq 10^9
  • 1\leq Q\leq 10^5
  • 1\leq V_i\leq N
  • 1\leq K_i\leq 20
  • 頂点V_iの部分木は頂点をK_i個以上持つ
  • 入力はすべて整数

解説

愚直(?)解

標準ライブラリのデータ構造でできる..はず。

K_iが小さいので、部分木を構成してしまえば楽勝。
vの部分木に書かれた整数を降順に数列にまとめたものをS_vとする。
S_v自体は、std::multisetとかを使えばいい。
このとき、vとvの子c_1,c_2,\dots,c_kについては

S_v=\{X_v\}\cup\left(\sum{S_{c_i}}\right)

が成り立つ。これを使うと、DFSとかで帰りがけ順に探索すればS_vを順に構成できる。

実装例(TLE)

#include<bits/stdc++.h>
using namespace std;
using graph = vector<vector<int>>;
int main() {
	int n, q;
	cin >> n >> q;
	vector<int> x(n);
	for (int i = 0; i < n; i++) {
		cin >> x[i];
	}
	graph g(n);
	for (int i = 0; i < n - 1; i++) {
		int a, b;
		cin >> a >> b;
		a--;
		b--;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	vector<multiset<int>> S(n);
	auto dfs = [&](auto f, int v, int par = -1)-> void {
		S[v].insert(x[v]);    //x[v]を追加する
		for (auto c : g[v]) {
			if (c == par)continue; //親ならcontinue
			f(f, c, v);   //S[v]を構築する
			//S[v] <- S[v] ∪ S[c]
			for (auto el : S[c]) {
				S[v].insert(el);
			}
		}
	};
	dfs(dfs, 0);
	while (q--) {
		int v, k;
		cin >> v >> k;
		v--;
		auto itr = S[v].end();
		for (int i = 0; i < k; i++) {
			itr--;
		}
		cout << (*itr) << endl;
	}
}

高速化

これだとTLEする。なぜなら、

S_v\leftarrow S_v + S_{c_i}

と更新するとき、一直線とかになるとO(N^2)かかるから。

具体例

上のようなグラフについて、

  • S_4を構成するとき、|S_5|=1より1回
  • S_3を構成するとき、|S_4|=2より2回
  • S_2を構成するとき...

...
みたいになって、計算量が1+2+\dots+N=\frac{N(N+1)}{2}(=O(N^2))です。

S_vをもう少しうまくやろう。



...再びK_i\leq 20の制約に注目する。
S_{V_i}について、K_i番目より後の要素のことはどうでもいい。もっと言うと、K_i\leq 20より全てのS_vについて20番目より後の要素のことはどうでもいい。
ということで、S_vの構成の際に

S_v\leftarrow S_v + S_{c_i}

としたあと、S_vのサイズが20以下になるまで小さいものを順に要素を除いていく。

実装例(AC)

計算量は...定数倍が20ぐらいのO(N)ですかね。

#include<bits/stdc++.h>
using namespace std;
using graph = vector<vector<int>>;

int main() {
	int n, q;
	cin >> n >> q;
	vector<int> x(n);
	for (int i = 0; i < n; i++) {
		cin >> x[i];
	}
	graph g(n);
	for (int i = 0; i < n - 1; i++) {
		int a, b;
		cin >> a >> b;
		a--;
		b--;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	vector<multiset<int>> S(n);
	auto dfs = [&](auto f, int v, int par = -1)-> void {
		S[v].insert(x[v]);
		for (auto c : g[v]) {
			if (c == par)continue;
			f(f, c, v);
			for (auto el : S[c]) {
				S[v].insert(el);
			}
		}
		//ここを追加した
		while (S[v].size() > 20) {
			S[v].erase(S[v].begin());
		}
	};
	dfs(dfs, 0);
	while (q--) {
		int v, k;
		cin >> v >> k;
		v--;
		auto it = S[v].end();
		for (int i = 0; i < k; i++) {
			it--;
		}
		cout << (*it) << endl;
	}
}

あんまり賢いことしなくても良いという話でした。

Discussion