問題概要

正の整数からなる長さ $N$ の数列 $A_{1}, A_{2}, \dots, A_{N}$ がある。頂点 $1, 2, \dots, N$ をもつ木をすべて考えたときの、

$\sum_{i=1}^{N} A_{i} d_{i}^{2}$

の最小値を求めよ ( $d_{i}$ は頂点 $i$ の次数)。

制約

$2 \le N \le 2 \times 10^{5}$

考えたこと

まず、次のことが自明に成り立つ。

$d_{i} \ge 1$
$d_{1} + d_{2} + \dots + d_{N} = 2N - 2$

ここでよく知られたこととして、逆にこれを満たす $d_{1}, d_{2}, \dots, d_{N}$ に対して、次数がこの値になるような木が存在する。このことは、プルーファーコードを用いたり、葉が必ず存在することから帰納法を活用したりすることで、証明できる。

よって問題は次の最適化問題を解くことに帰着された。

Maximize:

$A_{1} x_{1}^{2} + A_{2} x_{2}^{2} + \dots + A_{N} x_{N}^{2}$

Subject to:

$x_{i} \ge 1$
$x_{i}$ は整数
$x_{1} + x_{2} + \dots + x_{N} = 2N-2$

これは資源配分問題と呼ばれるものになっている。 $2N-2$ 個という有限個の資源を用いて、凸関数の和を最小化する資源配分問題だ。このような問題では、

$(x_{1}, x_{2}, \dots, x_{N}) = (1, 1, \dots, 1)$

から出発して、「 $x_{i}$ の値を 1 上昇させたときの目的関数の増分が最小であるような $i$ に対して、[tex: x_{i} の値を 1 上昇させる」ことを繰り返す Greedy 解法が最適解を導く。

たとえば priority_queue で実装できて、その場合の計算量は $O(N \log N)$ となる。

コード

#include <bits/stdc++.h>
using namespace std;
using pll = pair<long long, long long>;

int main() {
    long long res = 0;
    int N;
    cin >> N;
    vector<long long> A(N), x(N, 1);
    priority_queue<pll, vector<pll>, greater<pll>> que;
    for (int i = 0; i < N; ++i) {
        cin >> A[i], res += A[i];
        que.push(pll(A[i] * 3, i));
    }
    
    for (int iter = 0; iter < N - 2; ++iter) {
        auto [add, id] = que.top();
        que.pop();
        ++x[id];
        res += add;
        que.push(pll(A[id] * (x[id] * 2 + 1), id));
    }
    cout << res << endl;
}