RT,楼主每次提交都WA on test 5 ...
简单说一下思路:dfs得到每个节点的子树节点数,然后算出这条边的weight,即subtree[u]⋅(N−subtree[u]). 然后分类讨论一下M 和N−1的大小,贪心地选取权值。
long long该开的应该也都开了...请大佬们帮帮萌新吧,谢谢啦!
代码如下,码风还行:
const int mod = 1e9 + 7;
const int maxn = 1e5 + 5;
const int maxm = 6e4 + 5;
int subtree[maxn], p[maxm], N, M, T, cnt;
ll weight[maxn];
vector<int> G[maxn];
void dfs(int u, int k) {
subtree[u] = 1;
for (int i = 0; i < (int) G[u].size(); i++) {
int v = G[u][i];
if (v == k) continue;
dfs(v, u);
subtree[u] += subtree[v];
}
if (u != 1) weight[++cnt] = (ll) subtree[u] * (ll) (N - subtree[u]) % mod;
}
void solve() {
ll ans = 0;
cin >> N;
rep(i, 1, N) G[i].clear();
memset(weight, 0, sizeof(weight));
cnt = 0;
rep(i, 1, N - 1) {
int u, v;
cin >> u >> v;
G[u].pb(v), G[v].pb(u);
}
dfs(1, -1);
cin >> M;
rep(i, 1, M) cin >> p[i];
sort(p + 1, p + M + 1, greater<int>());
sort(weight + 1, weight + N, greater<int>());
if (M <= N - 1) {
for (int i = 1; i <= M; i++) ans = ((ll) ans + (ll) weight[i] * p[i]) % mod;
for (int i = M + 1; i <= N - 1; i++) ans = ((ll) ans + (ll) weight[i]) % mod;
} else {
ll mul = 1;
rep(i, 1, M - N + 2) mul = (ll) mul * (ll) p[i] % mod;
ans = (ll) weight[1] * (ll) mul % mod;
rep(i, 2, N - 1) ans = ((ll) ans + (ll) weight[i] * p[M - N + 1 + i]) % mod;
}
cout << ans << endl;
}