[ARC139D] Priority Queue 2 题解

上个世纪做过这题，然后今天比赛（abc295）出了道弱化没做出来，被 pty 喷了一遍后爬来写个题解/kk

首先这种期望/总和题都有个套路，就是通过另外一种角度来计算每个元素的贡献。对于此题，我们有：

\[ans=\sum_{i=1}^mi\cdot c(=i)=\sum_{i=1}^mc(\ge i) \]

其中 \(c(P)\) 表示对于所有方案中的每个 \(a_i\)，满足条件 \(P\) 的数的数量。

于是考虑枚举 \(i\)，则我们需要维护 \(c(\ge i)\)，对一次操作，设这次操作选出的数为 \(v\)，我们有：

• \(v<i\)：\(c\) 不变。
• \(v\ge i\)：\(c\) 加一。

而在每次操作后，由于我们还要删除第 \(x\) 小的元素，所以还有：

• 操作完后，若 \(n+1-c<x\)，即 \(c>n+1-x\)，那么 \(c\) 减一。

可以发现，\(n+1-x\) 像一个界限，\(c\) 如果超过了它就会被压回去，所以我们可以先忽略删除第 \(x\) 小，在最后将 \(c\) 和 \(n+1-x\) 取 \(\min\) 即可。

考虑对于 \(k\) 次操作，其中有多少次操作将 \(c\) 加了一，设加了 \(p\) 次，那么方案数显然为 \((m-i+1)^p(i-1)^{k-p}\displaystyle{k\choose p}\)（没加一的方案数 乘上 加一的方案数 乘上 \(k\) 次操作中选出 \(p\) 次操作的方案数）。

由于最后的 \(c\) 可以直接计算得到，故直接累加答案即可。

一个细节：注意到减一最多减 \(k\) 次，所以还要定一个下界。

时间复杂度 \(O(m(n+k))\)。

#include <atcoder/all>
#include <iostream>

using namespace std;
using LL = atcoder::modint998244353;

const int kN = 2001;
const int kM = 998244353;

int n, m, k, x, a[kN];
LL ans, sc[kN], iv[kN];

int main() {
  ios_base::sync_with_stdio(0), cin.tie(0);
  cin >> n >> m >> k >> x;
  iv[1] = 1;
  for (int i = 2; i <= k; ++i) {
    iv[i] = -LL(kM / i) * iv[kM % i];
  }
  sc[0] = 1;
  for (int p = 1; p <= k; ++p) {
    sc[p] = sc[p - 1] * (k - p + 1) * iv[p];
  }
  for (int i = 1; i <= n; ++i) {
    cin >> a[i];
  }
  for (int i = 1; i <= m; ++i) {
    int bc = 0;
    for (int j = 1; j <= n; ++j) {
      bc += a[j] >= i;
    }
    for (int p = 0; p <= k; ++p) {
      int c = max(bc + p - k, min(bc + p, n + 1 - x));
      ans += LL(i - 1).pow(k - p) * LL(m - i + 1).pow(p) * sc[p] * c;
    }
  }
  cout << ans.val();
  return 0;
}