数据结构学傻了

题目

题目描述

x d m \rm xdm xdm 市供水系统是树状结构， 1 1 1 号水塔是总站（顶层水塔），上层水塔逐级输水至下层水塔。每个水塔都有一个蓄水量。

现在，在根节点 1 1 1 出现了一个口渴的 H a n d I n D e v i l \sf HandInDevil HandInDevil 。由于一些众所周知的原因，他每次可以选择一个叶节点，喝光叶节点到自己所在节点这条简单路径上的所有水！

显然一个水塔的水一旦被喝掉，就不会重新获得。现在 H a n d I n D e v i l \sf HandInDevil HandInDevil 准备喝 k k k 次水，请问他最多能喝到多少水？

数据范围与提示

水塔数量 n ⩽ 4 × 1 0 6 n\leqslant 4\times10^6 n⩽4×106，蓄水量是小于 1 0 9 + 7 10^9+7 109+7 的自然数。

由于供水系统可以任意修建，所以实际上它们满足这样一个规律：水塔的蓄水量依次为 A 1 , A 2 , … , A n A_1,A_2,\dots,A_n A1​,A2​,…,An​，则 i i i 号水塔的父节点是 ( A i ⊕ B )   m o d   ( i − 1 ) + 1 (A_i\oplus B)\bmod(i-1)+1 (Ai​⊕B)mod(i−1)+1，且蓄水量满足 A i = ( B A i − 1 + C )   m o d   M A_i=(BA_{i-1}+C)\bmod M Ai​=(BAi−1​+C)modM，其中 { B = 23333333 C = 6666666 M = 1 0 9 + 7 \begin{cases}B=23333333\\C=6666666\\M=10^9+7\end{cases} ⎩⎪⎨⎪⎧​B=23333333C=6666666M=109+7​ 。

H a n d I n D e v i l \sf HandInDevil HandInDevil 只会告诉你 n , k n,k n,k 以及他所在节点的蓄水量。你能帮帮他吗？

思路

显然可以每次选当前权值最大的链。否则，找到这条链自底向上遇到的第一个已经被喝掉的水塔，称为 x x x，将经过 x x x 的链换成当前链。由于当前链是全局最大，所以除去 x x x 以上的部分，二者相比较，肯定是当前链更大。

直接进行这个过程肯定不行。考虑 d f s \tt dfs dfs 时维护每一次出发获得的权值。将儿子子树全部合并后，肯定会先走权值最大的，所以将这个最大权值加上当前权值。可并堆维护一下就是 O ( n log ⁡ n ) \mathcal O(n\log n) O(nlogn) 的。

想来没有优于 log ⁡ n \log n logn 的排序算法吧（除了基数排序）？然而 真的需要排序吗？上面可并堆的过程中，其实只对最大值进行了修改。如果我们只存最大值呢？

只存储最大值，只会导致剩下的值无序而已，本身过程是没有问题的。那么我们就得到了一个无序的数组，求出前 k k k 大的数的和即可。这就是 n t h _ e l e m e n t \rm nth\_element nth_element 啊！

不知道 S T L \rm STL STL 是怎么实现它的，反正效果是：第 k k k 小的元素在第 k k k 位，且它左边都比它小、右边都比它大。我还傻乎乎地手写了一个，真该像 沐目女未 一样扪心自问一下。

时间复杂度是期望 O ( n ) \mathcal O(n) O(n) 的！对这个程序的主要优化点就是 “前 k k k 大的和” 有线性做法。

代码

如果手写 n t h _ e l e m e n t \rm nth\_element nth_element，要考虑当前序列的所有元素都相同，就必须把 m i d \rm mid mid 单独计数。以及递归非常慢，请改成循环！

#include <cstdio>
#include <iostream>
#include <cstring>
#include <vector>
#include <algorithm>
#include <queue>
using namespace std;
# define rep(i,a,b) for(int i=(a); i<=(b); ++i)
# define drep(i,a,b) for(int i=(a); i>=(b); --i)
typedef unsigned long long int_;
inline int readint(){
	int a = 0; char c = getchar(), f = 1;
	for(; c<'0'||c>'9'; c=getchar())
		if(c == '-') f = -f;
	for(; '0'<=c&&c<='9'; c=getchar())
		a = (a<<3)+(a<<1)+(c^48);
	return a*f;
}
inline void writeint(int x){
	if(x > 9) writeint(x/10);
	putchar((x-x/10*10)^48);
}

const int MaxN = 4000005;
unsigned son[MaxN], fa[MaxN];
int_ val[MaxN];

int_ __builtin_median(int_ a,int_ b, int_ c){
	if(a <= b && a <= c) return min(b,c);
	if(b < a && b <= c) return min(a,c);
	/* if(c < a && c < b) */ return min(a,b);
}
int_ first_k_sum(int_ *l,int_ *r,unsigned k){
	if(!k || l == r) return 0;
	if(r-l == 1) return *l; // with k >= 1
	int_ *i = l+((r-l)>>1), *j = r-1;
	int_ mid = __builtin_median(*l,*i,*j);
	unsigned siz = 0; // count of %mid%
	for(i=l; i!=r; ++i){
		if((*i) != mid) continue;
		swap(*i,*l), ++ l, ++ siz;
	}
	for(i=l; i!=j+1; ){
		while(i != r && (*i) < mid) ++ i;
		while(j != i-1 && mid < (*j)) -- j;
		if(i != j+1) swap(*i,*j), ++ i, -- j;
	}
	if(k < r-i) // all in right part
		return first_k_sum(i,r,k);
	int_ res = 0; k -= (r-i);
	for(j=i; j!=r; ++j) res += (*j);
	if(k <= siz) return res+k*mid;
	res += siz*mid; k -= siz;
	return res+first_k_sum(l,i,k);
}

const unsigned BASE = 23333333;
const unsigned CS = 6666666;
const unsigned Mod = 1e9+7;
int main(){
	int n = readint(), k = readint();
	val[1] = readint(); son[1] = 1;
	rep(i,2,n){
		val[i] = (BASE*val[i-1]+CS)%Mod;
		fa[i] = (val[i]^BASE)%(i-1)+1;
		son[i] = son[fa[i]] = i;
	}
	for(unsigned i=n; i; --i){
		if(son[i] != i){ // not leaf
			val[son[i]] += val[i];
			val[i] = 0; // clear
		}
		if(val[son[fa[i]]] < val[son[i]])
			son[fa[i]] = son[i];
	}
	printf("%llu\n",first_k_sum(val+1,val+n+1,k));
	return 0;
}
           

彩蛋：有没有人发现题目中的 x d m \rm xdm xdm 是 “西多摩市” 的意思（滑稽）