Codeforces Round #665 (Div. 2) D. Maximum Distributed Tree

题意

给定一棵\(n\) 个结点的树，对这棵树分配边权，使得这棵树的边权的乘积为\(k\) ，且要求所有两点的简单路径边权之和最大。

\(k\) 以质因子的形式给出，有\(m\) 个质因子

结果取余1e9+7

\[n \leq 10^5 ,m \leq 6\cdot 10^4 ,p_i \leq 6\cdot 10^4 \]

分析

显然我们需要考虑每条边的贡献，如何求得一条边的贡献次数？简单排列组合一下就知道，令\(siz(i)\) 为\(i\) 号结点的子树大小，那么\(i\) 和父亲结点的连边的贡献次数就是$ size(i) \cdot (n - siz(i))$ ，这样我们只要贪心地把最大的质因子分配给贡献最大的连边即可，至于为什么是质因子，可以通过不等式证明。

这里还需要简单讨论一下 $ m \leq n - 1$ 的情况。

代码

ll siz[maxn], p[maxn], q[maxn];
vector<int> e[maxn];

void dfs(int u, int fa) {
    siz[u] = 1;
    for (auto it : e[u]) {
        if (it == fa) continue;
        dfs(it, u);
        siz[u] += siz[it];
    }
}

int main() {
    int T = readint();
    while (T--) {
        int n = readint();
        for (int i = 1; i <= n; i++) e[i].clear(), siz[i] = 0, p[i] = 1;
        for (int i = 1; i < n; i++) {
            int x, y;
            scanf("%d%d", &x, &y);
            e[x].push_back(y);
            e[y].push_back(x);
        }
        dfs(1, 0);
        int m = readint();
        for (int i = 1; i <= m; i++) p[i] = readll();
        if (m < n - 1) m = n - 1;
        sort(p + 1, p + m + 1);
        for (int i = n; i <= m; i++) {
            p[n - 1] *= p[i];
            p[n - 1] %= MOD;
        }
        ll res = 0;
        for (int i = 1; i < n; i++)
            q[i] = (n - siz[i + 1]) * siz[i + 1];
        sort(q + 1, q + n);
        for (int i = 1; i < n; i++)
            res = (res + (p[i] % MOD * q[i] % MOD) % MOD) % MOD;
        Put(res);
        puts("");
    }
}