bzoj 1267 Kth Number I (点分治,堆)

超级钢琴的树上版本, 类似做法即可, 只不过区间转为dfs序了, 用点分求一下, 复杂度$O(nlog^2n)$

#include <iostream>
#include <algorithm>
#include <cstdio>
#include <math.h>
#include <set>
#include <map>
#include <queue>
#include <string>
#include <string.h>
#include <bitset>
#define REP(i,a,n) for(int i=a;i<=n;++i)
#define PER(i,a,n) for(int i=n;i>=a;--i)
#define hr putchar(10)
#define pb push_back
#define lc (o<<1)
#define rc (lc|1)
#define mid ((l+r)>>1)
#define ls lc,l,mid
#define rs rc,mid+1,r
#define x first
#define y second
#define io std::ios::sync_with_stdio(false)
#define endl '\n'
using namespace std;
typedef long long ll;
typedef pair<int,int> pii;
const int P = 1e9+7, INF = 0x3f3f3f3f;
ll gcd(ll a,ll b) {return b?gcd(b,a%b):a;}
ll qpow(ll a,ll n) {ll r=1%P;for (a%=P;n;a=a*a%P,n>>=1)if(n&1)r=r*a%P;return r;}
ll inv(ll x){return x<=1?1:inv(P%x)*(P-P/x)%P;}
//head




const int N = 1e6+10;
int n, m, sum, rt;
struct _ {int to,w;};
vector<_> g[N];
int mx[N], sz[N], vis[N];
#define go for (int i=0,y,w; i<g[x].size()?w=g[x][i].w,y=g[x][i].to:0; ++i) 
void getrt(int x, int fa) {
	mx[x]=0,sz[x]=1;
	go if (!vis[y]&&y!=fa) {
		getrt(y,x),sz[x]+=sz[y];
		mx[x]=max(mx[x],sz[y]);
	}
	mx[x]=max(mx[x],sum-sz[x]);
	if (mx[rt]>mx[x]) rt=x;
}

int s[N], L, R;
pii a[N];

void dfs(int x, int f, int d, int rt) {
	s[++*s]=d, a[*s]=pii(L,R);
	go if (!vis[y]&&y!=f) dfs(y,x,d+w,rt);
}

void solve(int x) {
	vis[x]=1,s[++*s]=0,L=R=*s;
	go if (!vis[y]) dfs(y,x,w,x),R=*s;
	go if (!vis[y]) {
		mx[rt=0]=n,sum=sz[y];
		getrt(y,0),solve(rt);
	}
}

int Log[N], f[N][20];
void init(int n) {
	Log[0]=-1;
	REP(i,1,n) f[i][0]=i,Log[i]=Log[i>>1]+1;
	for (int j=1; (1<<j)<=n; ++j) {
		for (int i=0; i+(1<<j)-1<=n; ++i) {
			int x=f[i][j-1], y=f[i+(1<<(j-1))][j-1];
			f[i][j]=s[x]>s[y]?x:y;
		}
	}
}
int RMQ(int l, int r) {
	int k=Log[r-l+1];
	int x=f[l][k],y=f[r-(1<<k)+1][k];
	return s[x]>s[y]?x:y;
}
struct node {
	int l,r,pos,opt,v;
	node () {}
	node (int l, int r, int pos) : l(l),r(r),pos(pos) {
		opt = RMQ(l,r);
		v = s[pos]+s[opt];
	}
	bool operator < (const node & rhs) const {
		return v < rhs.v;
	}
};

priority_queue<node> q;
int main() {
	scanf("%d%d", &n, &m);
	REP(i,2,n) {
		int u, v, w;
		scanf("%d%d%d", &u, &v, &w);
		g[u].pb({v,w}),g[v].pb({u,w});
	}
	sum=mx[0]=n,getrt(1,0),solve(rt);
	init(*s);
	REP(i,1,*s) if (a[i].x) q.push(node(a[i].x,a[i].y,i));
	REP(i,1,m) {
		node t = q.top();
		printf("%d\n", t.v);q.pop();
		if (t.opt!=t.l) q.push(node(t.l,t.opt-1,t.pos));
		if (t.opt!=t.r) q.push(node(t.opt+1,t.r,t.pos));
	}
}