資料結構學傻了

題目

題目描述

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 是 “西多摩市” 的意思（滑稽）