【题解】ARC139D Priority Queue 2
？
题意
假设初始时有一个长度为 \(N\)，值域为 \(M\) 的数组 \(A\)。现在要进行 \(K\) 次操作，每次操作从 \([1, M]\) 中选取一个数，并将其加入 \(A\) 中。单次操作完成后对 \(A\) 按升序排序，然后将第 \(X\) 大的元素删除。对于 \(M^K\) 种可能的操作序列，求 \(\sum\limits_{i = 1}^n A_i\) 的和。答案对 \(998244353\) 取模。
\(N, M, K \leq 2 \times 10^3, 1 \leq X \leq N + 1\)
思路
带玄学神仙计数。
首先会发现操作实际上等价于每次不删除第 \(X\) 小的元素，然后对于最终得到的长度为 \(N + K\) 的序列保留：
-
前 \(X - 1\) 小的数
-
后 \(N - X + 1\) 大的数
并且求这两种情况的贡献是等价的问题，于是我们考虑解决其中一种，另一种类比即可。这里以求前 \(X - 1\) 小的贡献为例。
假设最终序列中第 \(X\) 小的数为 \(v\)，初始时 \(A\) 中小于 \(v\) 的元素有 \(a\) 个，大于 \(v\) 的元素有 \(b\) 个。那么在加入的 \(K\) 个数中：
-
小于 \(v\) 的数不超过 \((X - 1) - a\) 个
-
大于 \(v\) 的数不超过 \(K - (X - a - b)\) 个，即 \(K + a + b - X\) 个
不妨枚举 \(i, j\) 分别表示小于 \(v\) 的数的个数和大于的个数。考虑当加入的 \(K\) 个数中有 \(i\) 个数小于 \(v\) 时，大于 \(v\) 的数的取值方法 \(C_i\)：
\(C_i = \sum\limits_{j = 0}^{K + a + b - X} C_{K - i}^{j} \cdot (M - v)^{j}\)
注意到上式的变化仅与 \(i\) 有关，于是可以 \(O(K)\) 求。
那么前 \(X - 1\) 个数的贡献为：
-
小于 \(v\) 的数，贡献之和为 \(\sum\limits_{i = 0}^{X - 1 - a} C_K^i \cdot C_i \cdot f_{v - 1, i}\)
其中 \(f_{v - 1, i} = (v - 1)^{i - 1} \cdot \frac{v(v - 1)}{2}\)，即第 \(i\) 位取 \([1, m]\) 时对答案的贡献之和。 -
等于 \(v\) 的数，贡献之和为 \(\sum\limits_{i = 0}^{X - 1 - a} C_K^i \cdot C_i \cdot [v \cdot (X - 1 - i)]\)
枚举 \(v\) 求贡献和即可。
后 \(N - X + 1\) 大的数的贡献同理。
时间复杂度 \(O(MK)\)
