買うサイコロを決め打った時の総和の二乗の期待値を単純な式で表してみます。添え字を簡単にするために、サイコロ \(1\) から \(K\) を買うとしましょう。

ここで確率母関数 \(f_i(x) =\frac{1}{6}(\sum_{j=1}^{6} x^{A_{i,j} }) \) を導入します。また \(F(x) \coloneqq \prod_{i=1}^K f_i(x)\) とすると、総和が \(k\) になる確率は \([x^k]F(x)\) で得られます。

二乗の期待値はどのように求められるか考えると、\(G(x) \coloneqq (x F'(x))'\) としたときの \(G(1)\) に一致していることが分かります。

積の微分法則を用いて \(G(1)\) を求めましょう。\(A_i = f_i'(1), B_i = f_i''(1)\) と置きます。このとき

\[G(1) = \sum_{i=1}^K A_i +2 \sum_{i =1}^{K-1}\sum_{j=i+1}^{K} A_i A_j + \sum_{i=1}^K B_i\]

であり、式変形して

\[G(1) = \sum_{i=1}^K A_i +\big( \sum_{i=1}^K A_i\big) ^ 2 - \sum_{i=1}^K A_i^2 + \sum_{i=1}^K B_i\]

、すなわち

\[G(1) = \big( \sum_{i=1}^K A_i\big) ^ 2+ \sum_{i=1}^K (A_i + B_i - A_i ^2) \]

が得られました。

以上の考察により、\(x_i = A_i, y_i = A_i + B_i - A_i ^2 - C_i \) とすると本問題は以下の問題に言い換えられます。

---

\(N\) 個のベクトル \((x_i,y_i)\) が与えられる。これらのベクトルの中からちょうど \(K\) 個を選び、選んだベクトルの \((\sum x) ^ 2 + \sum y \) を最大化せよ。

---

この問題は以下の手順で解くことができます。

最適解の構造を分析します。すると、最適解に対してある実数 \(c\) が存在し、\(cx +y\) が大きい順に \(K\) 個取っていることがわかります。

---

補足

平面上の線分 \(PQ\) を考えます。この線分上を点 \(R\) が等速に動くとき (x 座標)^2 + (y 座標) の値は下に凸な二次関数となるので、最大値は端点のどちらかでとります。 このことを利用すると、一般に、平面上の点集合が与えられたとき、\(x^2+y\) を最大化する点はその点集合の凸包の頂点をなす点のみ調べればよいことが分かります。

\(N\) 個のベクトルから \(K\) 個選んだ和として考えられるベクトル全てに対応する \(\binom{N}{K}\) 点を平面上にプロットすることを考えましょう。この点集合の凸包の頂点を列挙できればよいです。(上側) 凸包の頂点は、ある実数 \(c\) に対して \(cx+y\) を最大化する点なので、この \(c\) を固定して \(\binom{N}{K}\) 点のうち \(cx+y\) を最大化する点を探すことを考えます。これは元の \(N\) 点の中から \(cx+y\) の大きい順に \(K\) 点をとればよいです。

---

また、\(c\) の候補としては \(2\) 点 \((x_i,y_i),(x_j,y_j)\) を通る直線の傾きのみ考えれば良いです。

直線は \(O(N^2)\) 個であり、\(c\) を決め打つと \(O(N \log N)\) で貪欲に取れるので、 \(O(N^3\log N)\) でこの問題に答えることが出来ますが、これでは実行時間に間に合わせるのは困難です。

先述の解法を偏角ソートにより \(O(N^2\log N)\) に改良しましょう。

\(x\) 座標の降順にベクトルをソートしておきます。また、上位 \(K\) 個の \(x\) の総和と \(y\) の総和も事前に求めておきましょう。

\(O(N^2)\) 個の直線を偏角ソートし、順に見ていきます。ある直線を見るときはその直線を通る \(2\) 点のベクトルを交換し、上位 \(K\) 個の \(x\) の総和と \(y\) の総和を適切に更新します。

この工夫により計算量を \(O(N^2\log N)\) に落とすことができます。これは十分高速です。