AtCoder ABC239 E - Subtree K-th Max

题目描述

给定一棵 \(n\) 个节点的树，\(i\) 节点的权值为 \(x_i\)，根节点编号为 \(1\)。

现有 \(Q\)个询问，每个询问给定 \(v,k\)，求节点 \(v\) 的子树第 \(k\) 大的数。

输入输出样例

4
5

数据范围

\(0 \le x_i \le 10^9\)

\(2 \le n \le 10^5\)

\(1 \le Q \le 10^5\)

\(1 \le k \le 20\)

算法

DFS，树形结构

分析

如果对于每次询问都去访问子树，显然时间超限。

观察数据范围，会发现 \(1 \le k \le 20\)，所以可以预处理出所有节点的子树的前 \(20\) 大。

处理大小，可以使用一个大顶堆来维护。

规定，\(ans_{i, j}\) 表示 \(i\) 点的子树中的第 \(j\) 大是多少。

那么，如何找到一棵子树的前 \(20\) 大呢？

如果要寻找 \(2\) 这个节点的前 \(20\) 大，那么就需要在它的儿子节点的子树中找到前 \(40\) 大，在从中找到前 \(20\) 大。因此对于每一个节点，首先需要遍历它的儿子节点，在这些节点中递归地寻找前 \(20\) 大的节点，并更新所有的 \(ans_{i, j}\)。

DFS 过程结束后，对于每次询问，直接输出 \(ans_{v, k}\) 即可。

代码

#include <iostream>
#include <queue>

using namespace std;

#define int long long

const int N = 1e5 + 10;

// w[i] 表示第 i 个点的边权 
int n, q, w[N], a, b, v, k;
vector<int> g[N];
int ans[N][30];

priority_queue<int> heap, c; 

// 深度优先遍历 
void dfs(int u, int fa){
	ans[u][1] = w[u];
	
	for(auto v : g[u]){
		// 如果 v 点不是 u 的儿子节点，不做处理 
		if(v == fa){
			continue;
		}
		
		dfs(v, u);	// v 点的 ans 我们认为已经求解完毕 
		heap = c;	// 初始化 
		
		// 将原来已经存储好的 ans[u][i] 和现在的 ans[v][i] 进行合并，并找到前 20 大 
		for(int i=1; i<=20; i++){
			heap.push(ans[u][i]); 
			heap.push(ans[v][i]);
		}
		for(int i=1; i<=20; i++){
			ans[u][i] = heap.top();
			heap.pop();
		}
	}
}

signed main(){
	// 读入 
	scanf("%lld %lld", &n, &q);
	
	for(int i=1; i<=n; i++){
		scanf("%lld", w+i); 
	}
	
	for(int i=1; i<n; i++){
		scanf("%lld %lld", &a, &b);
		// 建树 
		g[a].push_back(b);
		g[b].push_back(a);
	}
	
	// 遍历 
	dfs(1, 0);
	
	// 访问并输出 
	while(q--){
		scanf("%lld %lld", &v, &k);
		printf("%lld\n", ans[v][k]);
	}
	
	return 0;
}