[ural 2120]. Tree Average Weight

题意

给出一个长度为\(n\)的序列\(\{a_i\}\)，定义一棵\(n\)个点的树为好树当且仅当对于每个节点有：

1. 若\(a_i = -1\)，\(i\)号节点度数随意；
2. 若\(a_i \neq -1\)，\(i\)号节点度数为\(a_i\)；

定义一棵树的权值为\(\sum_{(u, v) \in E} u * {sz}_{uv}(u) + v * {sz}_{vu}(v)\)。
其中\({sz}_{uv}(u)\)表示删去边\((u, v)\)后，\(u\)所在的连通块的大小。
求所有好树的平均值的整数部分。
保证至少存在一棵好树。
\(n \leq 1e6\)。

题解

考虑一棵好树的权值为

\[\begin{aligned} & \sum_{x} \sum_{y} [x, y \ is \ adjacent] x * (n - {sz}_y) \\ = & \sum_{x} \sum_{y} [x, y \ is \ adjacent] x * (n - {sz}_y) \\ = & \sum_{x} x * (n * \deg(x) - (n - 1)) \\ = & \sum_{a_x > 0} x * (n * a_x - (n - 1)) + \sum_{a_x = -1} x * (n * \deg(x) - (n - 1)) \\ \end{aligned} \]

我们对于所有好树一起算权值。
前一部分显然可以直接计算，我们设为\(former\)；考虑后一部分，枚举度数进行计算。
设\(s = \sum_{a_x > 0} a_x\)，\(t = \sum_{a_x = -1} x\)，则后一部分为

\[\sum_{i = 1} ^ {n - s - 1} n * t * i * ways(i) \]

\(ways(i)\)即为钦定某一个\(a_x = -1\)的点度数为\(i\)，能形成的不同的好树数。
根据prufer序列，有

\[ways(i) = \frac{(n - 2)!}{D (n - 1 - s - i)!} * {(c - 1)} ^ {n - 1 - s - i} \]

其中\(c = \sum_{a_x = -1} 1\)，\(D = \prod_{a_x > 0} (a_x - 1)!\)。
看上去权值和部分已经计算好了，不妨来计算总的好树个数。
同样也是用prufer序列，有

\[\#goodtree = \frac{(n - 2)!}{D(n - 2 - s)!} * c ^ {n - 2 - s} \]

根据问题的输出格式，又没有取模，肯定要化简以后才能算。
答案为

\[former + n * t * \frac{\sum_{i = 1} ^ {n - s - 1} \frac{(n - 2)!}{D (n - 1 - s - i)!} * {(c - 1)} ^ {n - 1 - s - i}}{\frac{(n - 2)!}{D(n - 2 - s)!} * c ^ {n - 2 - s}} \]

令\(r = n - s - 1\)，则（经过一波化简得到）答案为

\[\begin{aligned} & former + n * t * \frac{1}{c ^ {r - 1}} \sum_{i = 1} ^ r \frac{i (r - 1)! {(c - 1)} ^ {r - i}}{(i - 1)! (r - i)!} \\ = & former + n * t * \frac{1}{c ^ {r - 1}} \sum_{i = 1} ^ r \frac{(i - 1 + 1) (r - 1)! {(c - 1)} ^ {r - i}}{(i - 1)! (r - i)!} \\ = & former + n * t * \frac{1}{c ^ {r - 1}} \left[\sum_{i = 1} ^ {r - 1} \binom{r - 1}{i - 1} {(c - 1)} ^ {r - 1 - i} + (r - 1) \sum_{i = 2} ^ {r - 2} \binom{r - 2}{i - 2} {(c - 1)} ^ {r - 2 - i} \right] \\ = & former + n * t * \frac{1}{c ^ {r - 1}} \left( (1 + c - 1) ^ {r - 1} + (r - 1) (1 + c - 1) ^ {r - 2} \right) \\ = & former + n * t * (1 + \frac{r - 1}{c}) \end{aligned} \]

现在问题就是求\(former + n * t * (1 + \frac{r - 1}{c})\)的整数部分。
也就是处理\(\lfloor \frac{a * b}{c} \rfloor\)的问题，其中\(a * b\)会爆long long，但\(c\)比较小，是\(1e6\)级别。
暴力拆成\(\frac{(Ac + x)(Bc + y)}{c}\)计算即可。
复杂度\(\mathcal O(n)\)。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int N = 1e6 + 5;
ll n, x, sum, cnt, tmp, ans1, ans2;
ll calc (ll x, ll y, ll z) {
	ll a = x / z, b = y / z;
	x %= z, y %= z;
	return a * b * z + a * y + b * x + x * y / z;
}
int main () {
	cin.tie(0), ios :: sync_with_stdio(0);
	cin >> n;
	for (int i = 1; i <= n; ++i) {
		cin >> x;
		if (x > 0) {
			sum += x - 1;
			ans1 += x * n * i;
		} else {
			++cnt;
			tmp += i;
		}
		ans1 -= (n - 1) * i;
	}
	sum = n - sum - 1;
	ans2 = cnt ? calc(tmp * n, sum + cnt - 1, cnt) : 0;
	cout << ans1 + ans2 << endl;
	return 0;
}