こちらはユーザ解説です。別解を紹介します。

期待値ではなく、全ての操作についての \(\prod_{i = 1}^N B_i\) の総和を求める方法を説明します。この値を \(N^K\) で割ることで答えが求められます。

まず、任意の \(1 \leq i \leq N\) に対し \(A_i = 0\) である場合について考えます。操作は \(1\) 以上 \(N\) 以下の整数からなる長さ \(K\) の数列 \(C_1, \dots, C_K\) と一対一に対応し、\(B_i\) は \(C\) の要素のうち \(i\) に等しいものの個数です。 また、\(\prod_{i = 1}^N B_i\) は、各 \(1 \leq i \leq N\) について、\(C_k = i\) となる要素をちょうど \(1\) つ選び、印をつける方法の総数に等しいです。よって、求める値は、以下の条件を満たす数列 \(D\) の総数に等しいです。

• \(D\) の長さは \(K\) である。
• \(D\) の要素には、\(-1, -2, \dots, -N\) がちょうど \(1\) つずつ含まれる。
• 上記以外の \(D\) の要素は、\(1, 2, \dots, N\) のいずれかであり、これらはそれぞれ何個含まれていてもよい。

「印をつける」という操作を、「\(-1\) 倍する」と読み替えると納得しやすいでしょう。前述の条件を満たす数列 \(D\) の総数は \(K \times (K - 1) \times \dots \times (K - N + 1) \times N^{K - N}\) と簡単に表せ、\(O(N)\) で求められます。

\(A_i = 0\) という条件を外し、もとの問題について考えましょう。先ほどと同様に、\(\prod_{i = 1}^N B_i\) を、それぞれの箱についてボールを一つずつ選んで印をつける方法の総数、と言い換えます。印をつけられるボールが元からあった \(A_i\) 個のボールのいずれかであるような箱の個数が \(n\) 個であるとします。この \(n\) 個の箱の選び方と、それらに含まれるボールの印の付け方の総数は動的計画法により \(O(N^2)\) で計算できます。残りの \(N - n\) 個の箱については、元からあったボールには印をつけないので、\(A_i = 0\) である場合に帰着できます。

以上から、この問題を \(O(N^2)\) で解くことができました。

実装例 (C++) :