LG10641

有一个显然的贪心，即每次都取剩余权值和最大的一条路径。由于每个点的权值都只计算一次，因此这样做能够确保前 \(k\) 次得到的答案最大。假如不取大的而先取小的，那么答案一定不会优于前者（大的可能没取到），从而保证了贪心的正确性。

接下来考虑如何维护最大值。每一次求出最大值后，同时记录最大值对应的点的编号。对于一个点 \(u\)，它的权值 \(w_u\) 被取走后将会导致以 \(u\) 为根的子树中每个点的总权值减少 \(w_u\)。子树修改？使用 dfs 序加线段树维护即可！

但是如果我们每次都暴力修改会超时。注意到每个权值只能用一次，也就是说只有一次修改是有效的！那么我们使用一个标记数组记录每个点是否被修改，遇到一个点已经被修改时，说明它的祖先也已经修改了（因为都是向根节点跳），直接退出即可。

总时间复杂度 \(O((n+k)\log n)\)。

#include <iostream>
#include <cstdio>
#include <vector>
#define int long long

using namespace std;

vector<int> G[200001];
int n,k,a[200001],w[200001],Lx[200001],Rx[200001],rev[200001],fa[200001],vis[200001],tot,ans;

struct T
{
	int mx,lzy,id;
}t[1000001];

void pushdown( int u )
{
	if( t[u].lzy )
	{
		t[u << 1].mx += t[u].lzy;
		t[u << 1].lzy += t[u].lzy;
		t[u << 1 | 1].mx += t[u].lzy;
		t[u << 1 | 1].lzy += t[u].lzy;
		t[u].lzy = 0;
	}
	return;
}

void pushup( int u )
{
	if( t[u << 1].mx >= t[u << 1 | 1].mx ) t[u].mx = t[u << 1].mx,t[u].id = t[u << 1].id;
	else t[u].mx = t[u << 1 | 1].mx,t[u].id = t[u << 1 | 1].id;
	return;
}

void dfs( int u , int f )
{
	fa[u] = f;
	Lx[u] = ++ tot;
	rev[tot] = u;
	a[u] += w[u];
	for( auto v : G[u] )
	{
		if( v == f ) continue;
		a[v] += a[u];
		dfs( v , u );
	}
	Rx[u] = tot;
}

void build( int u , int l , int r )
{
	if( l == r )
	{
		t[u].mx = a[rev[l]];
		t[u].id = rev[l];
		return;
	}
	int mid = ( l + r ) >> 1;
	build( u << 1 , l , mid );
	build( u << 1 | 1 , mid + 1 , r );
	pushup( u );
}

void update( int u , int l , int r , int L , int R , int x )
{
	if( L <= l && r <= R )
	{
		t[u].mx += x;
		t[u].lzy += x;
		return;
	}
	pushdown( u );
	int mid = ( l + r ) >> 1;
	if( L <= mid ) update( u << 1 , l , mid , L , R , x );
	if( R > mid ) update( u << 1 | 1 , mid + 1 , r , L , R , x );
	pushup( u );
}

T ma( T x , T y )
{
	if( x.mx >= y.mx ) return x;
	return y;
}

T query( int u , int l , int r , int L , int R )
{
	if( L <= l && r <= R )
		return t[u];
	pushdown( u );
	int mid = ( l + r ) >> 1;
	T res;
	res.mx = 0;
	if( L <= mid ) res = ma( res , query( u << 1 , l , mid , L , R ) );
	if( R > mid ) res = ma( res , query( u << 1 | 1 , mid + 1 , r , L , R ) );
	return res;
}

signed main()
{
	int u,v,nw;
	T tt;
	cin >> n >> k;
	for( int i = 1 ; i <= n ; i ++ )
		cin >> w[i];
	for( int i = 1 ; i < n ; i ++ )
	{
		cin >> u >> v;
		G[u].push_back( v );
		G[v].push_back( u );
	}
	dfs( 1 , -1 );
	build( 1 , 1 , n );
	for( int i = 1 ; i <= k ; i ++ )
	{
		tt = query( 1 , 1 , n , 1 , n );
		ans += tt.mx;
		nw = tt.id;
		while( nw != -1 )
		{
			if( vis[nw] ) break;
			vis[nw] = 1;
			update( 1 , 1 , n , Lx[nw] , Rx[nw] , -w[nw] );
			nw = fa[nw];
		}
	}
	cout << ans;
	return 0;
}