[ural 2120]. Tree Average Weight

Posted 2020-11-21 psimonw

题意

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

1. 若(a_i = -1)，(i)号节点度数随意；
2. 若(a_i eq -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)。

题解

考虑一棵好树的权值为
[ egin{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 > 0} x * (n * deg(x) - (n - 1)) \end{aligned} ]
我们对于所有好树一起算权值。
前一部分显然可以直接计算；考虑后一部分，枚举度数进行计算。
设(s = sum_{a_x > 0} a_x)，(t = sum_{a_x < 0} x)，则后一部分为
[ egin{aligned} sum_{i = 1} ^ {n - s - 1} 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;
}