\(A_C\) に \(1\) を加えておき、\(A_i\) が（高橋君がはじめに持っているものも含めた）色 \(i\) の靴下の枚数となるようにしておきます。適切に色の番号を付け替えることで、\(A_1 \leq A_2 \leq \ldots \leq A_N\) としてよいです。

まずは、高橋君がタンスに戻す靴下について考えます。高橋君が色 \(x\) と色 \(y\) の靴下を持っているとき、高橋君はタンスの中に多く存在する色の靴下をタンスの外に残しておくべきです。したがって、\(x < y\) のとき高橋君は色 \(x\) の靴下をタンスの中に戻し、色 \(y\) の靴下をタンスの外に残すとしてよいです（色 \(x\) の靴下と色 \(y\) の靴下が同じ数だけタンスの中に入っているときも添え字の大小でタイブレークを与えるということになります）。

この観察により、dp による計算が行えます。

色 \(i\) の靴下をタンスの外に持っているときに操作を終えるまでの靴下を取り出す回数の期待値を \(dp_i\) とします。次に取り出した靴下の色で場合分けすることで計算式を与えます。次に取り出した靴下が色 \(j\) であるとします。\(S \coloneqq \sum A_i - 1\)、すなわち \(S\) をタンスの中にある靴下の枚数とすると、以下が成り立ちます。

• \(i > j\) のとき、色 \(i\) の靴下をタンスの外に残すことになるのでこのときの寄与は \(\frac {A_jdp_i}{S}\) である。
• \(i = j\) のとき、操作は終了するのでこのときの寄与は \(0\) である。
• \(i < j\) のとき、 色 \(j\) の靴下をタンスの外に残すことになるのでこのときの寄与は \(\frac {A_jdp_j}{S}\) である。

したがって、\(dp_i = 1 + \displaystyle\sum_{j = 1}^{i-1} \frac{A_jdp_i}{S} + \sum_{j = i + 1}^{N} \frac{A_jdp_j}{S}\) となります。

移項をすることで \((1 - \displaystyle\sum_{j = 1}^{i-1} \frac{A_j}{S}) dp_i = 1 + \sum_{j = i + 1}^{N} \frac{A_jdp_j}{S}\) となるため \(i\) の降順にこの式に従って計算することで \(O(N^2)\) 時間で計算が行えますが、これでは十分に高速な計算となっていません。

しかし、計算のボトルネックとなっていた \(\displaystyle\sum_{j = 1}^{i-1} \frac{A_j}{S}\) および \(\displaystyle\sum_{j = i + 1}^{N} \frac{A_jdp_j}{S}\) は差分更新により計算が高速にできる形になっています。実際に、\(X, Y\) を \(dp_i\) の計算のときに使ったそれぞれの値として管理しておき、\(dp_{i - 1}\) を計算する前に \(X\) から \(\frac{A_{i - 1}}{S}\) を引き、\(Y\) に \(\frac{A_idp_i}{S}\) を足すなどの方法により計算できます。

これにより全体として \(O(N)\) 時間で dp の計算が行えました。法 \(P = 998244353\) に対する依存も含めると \(O(N \log P)\) や \(O(N + \log P)\) といった評価となります。

また、上の解法では \(A_i \leq 3000\) という制約は \(S\) が \(998244353\) の倍数にならないなど、適切に除算が行えることのみに利用されていますが、\(A_i\) が小さいことを利用して高速化を行うこともできます。

\(dp_i\) を同じ色の靴下が \(i\) 枚ある色の靴下をタンスの外に持っているときに操作を終えるまでの靴下を取り出す回数の期待値とすれば、dp の状態数が \(O(\max A_i)\) となり、上と同じような dp を差分更新を用いずに行っても \(O((\max A_i)^2)\) 時間となり十分高速です。