【AtCoder】Beginner Contest 380-G.Another Shuffle Window

题目链接

Problem Statement

You are given a permutation $P$ of $(1,2,\dots ,N)$ and an integer $K$.

Find the expected value, modulo $998244353$, of the inversion number of $P$ after performing the following operation:

• First, choose an integer $i$ uniformly at random between $1$ and $N-K+1$, inclusive.
• Then, shuffle $P_i,P_{i+1},\dots ,P_{i+K-1}$ uniformly at random.

What is the inversion number?

• The inversion number of a sequence $(A_1,A_2,\dots ,A_N)$ is the number of integer pairs $(i,j)$ satisfying $1\le i<j\le N$ and $A_i>A_j$.

What does “expected value modulo $998244353$” mean?

• It can be proved that the sought expected value is always rational. Under the constraints of this problem, when this value is represented as an irreducible fraction $\frac{P}{Q}$, it can also be proved that $Q\not\equiv 0(\mathrm{mod}998244353)$. Thus, there is a unique integer $R$ satisfying $R\times Q\equiv P(\mathrm{mod}998244353),0\le R<998244353$. Report this integer $R$.

中文题意

问题陈述

你得到一个 $(1,2,\dots ,N)$ 的排列 $P$ 和一个整数 $K$ 。

对 $P$ 的反转数求模 $998244353$ 的期望值：

—首先，在 $1$ 到 $N-K+1$ 之间均匀随机选择一个整数 $i$ 。
—然后，均匀随机洗牌 $P_i,P_{i+1},\dots ,P_{i+K-1}$ 。

反转数是多少？

• 序列 $(A_1,A_2,\dots ,A_N)$ 的反转数是满足 $1\le i<j\le N$ 和 $A_i>A_j$ 的整数对 $(i,j)$ 的个数。

“期望值模 $998244353$ ”是什么意思？

• 可以证明所寻求的期望值总是理性的。在本问题的约束下，当该值表示为不可约分数 $\frac{P}{Q}$ 时，也可以证明 $Q\not\equiv 0(\mathrm{mod}998244353)$ 。因此，有一个唯一的整数 $R$ 满足 $R\times Q\equiv P(\mathrm{mod}998244353),0\le R<998244353$ 。报告这个整数 $R$ 。

Constraints

• All input values are integers.
• $1\le K\le N\le 2\times 10^5$
• $P$ is a permutation of $(1,2,\dots ,N)$.

Input

The input is given from Standard Input in the following format:

$N$ $K$
$P_1$ $P_2$ $\dots$ $P_N$

Output

Print the answer in one line.

Sample Input 1

4 2
1 4 2 3

Sample Output 1

166374061

The operation changes the permutation $P$ into the following:

• $(1,4,2,3)$ … probability $1/2$
• $(4,1,2,3)$ … probability $1/6$
• $(1,2,4,3)$ … probability $1/6$
• $(1,4,3,2)$ … probability $1/6$

The expected value of the inversion number is $2\times \frac{1}{2}+3\times \frac{1}{6}+1\times \frac{1}{6}+3\times \frac{1}{6}=\frac{13}{6}$.

$\frac{13}{6}$ modulo $998244353$ is $166374061$, so print this number.

Sample Input 2

1 1
1

Sample Output 2

0

Sample Input 3

10 6
7 4 10 5 6 1 8 2 3 9

Sample Output 3

499122200

解法

当 $P_L,P_{L+1},\dots ,P_R$ 被洗牌时，反转数会发生怎样的变化？

• Block $X$ … $P_1,P_2,\dots ,P_{L-1}$
• Block $Y$ … $P_L,P_{L+1},\dots ,P_R$
• Block $Z$ … $P_{R+1},P_{R+2},\dots ,P_N$

考虑元素 $P_i$ 和 $P_j$ $i<j$ 在洗牌后的相对位置。

• 如果 $P_i$ 属于 $X$ 块， $P_j$ 属于 $X$ 块，那么洗牌后 $P_i$ 仍然比 $P_j$ 早。
• 如果 $P_i$ 属于 $X$ 块， $P_j$ 属于 $Y$ 块，那么洗牌后 $P_i$ 仍然比 $P_j$ 早。
• 如果 $P_i$ 属于 $X$ 块， $P_j$ 属于 $Z$ 块，那么洗牌后 $P_i$ 仍然比 $P_j$ 早。
• 如果 $P_i$ 属于 $Y$ 块， $P_j$ 属于 $Y$ 块……洗牌后， $P_i$ 以 $1/2$ 的概率早于 $P_j$ ， $P_j$ 以 $1/2$ 的概率早于 $P_i$ 。
• 如果 $P_i$ 属于 $Y$ 块， $P_j$ 属于 $Z$ 块，那么洗牌后 $P_i$ 仍然比 $P_j$ 早。
• 如果 $P_i$ 属于 $Z$ 块， $P_j$ 属于 $Z$ 块，那么洗牌后 $P_i$ 仍然比 $P_j$ 早。

因此，两个元素的相对位置保持不变，除非它们都在块Y中，而块Y被随机均匀地洗牌。
通过讨论可知，期望的反演数为（原 $P$ 的反演数） $-$ （原 $P_L,P_{L+1},\dots ,P_R$ $+$ $({(}_{R-L+1}{C}_2)/2)$ 的反演数）。

剩下的就是找到所有 $1\le i\le N-K+1$ 的 $P_i,P_{i+1},\dots ,P_{i+K-1}$ 的反转数。

这可以发现如下：

• 首先查找 $P_1,P_2,\dots ,P_K$ 的反转数，用 $I$ 表示。
• 对于 $i=K+1,K+2,\dots ,N$ ，重复如下操作：
• 从 $I$ 中减去 $P_{i-K+1},P_{i-K+2},\dots ,P_{i-1}$ 中小于 $P_{i-K}$ 的元素个数。
• 然后，将 $P_{i-K+1},P_{i-K+2},\dots ,P_{i-1}$ 中大于 $P_{i-K+1},P_{i-K+2},\dots ,P_{i-1}$ 的元素个数加到 $I$ 。
• 此时 $I$ 为 $P_{i-K+1},P_{i-K+2},\dots ,P_i$ 的反转数。

这可以通过处理段和查询来计算（使用像段树这样的数据结构）。

#include<bits/stdc++.h>
#include<atcoder/segtree>

using namespace std;
using namespace atcoder;

#define mod 998244353
#define inv2 499122177

long long op(long long a,long long b){
  if((a+b)>=mod){
    return (a+b)-mod;
  }
  return (a+b);
}
long long e(){
  return 0;
}

long long power(long long a,long long b){
  long long x=1,y=a;
  while(b>0){
    if(b&1ll){
      x=(x*y)%mod;
    }
    y=(y*y)%mod;
    b>>=1;
  }
  return x%mod;
}

long long modular_inverse(long long n){
  return power(n,mod-2);
}

int main(){
  ios::sync_with_stdio(false);
  cin.tie(nullptr);

  long long n,k;
  cin >> n >> k;
  vector<long long> p(n);
  for(auto &nx : p){
    cin >> nx;
    nx--;
  }

  long long bas=0;
  {
    segtree<long long,op,e> seg(n);
    for(long long i=0;i<n;i++){
      bas+=seg.prod(p[i],n);
      seg.set(p[i],1);
    }
  }

  long long res=0;
  long long sub=0;
  long long add=((k*(k-1))/2)%mod; add*=inv2; add%=mod;

  segtree<long long,op,e> seg(n);
  for(long long i=0;i<k;i++){
    sub+=seg.prod(p[i],n); sub%=mod;
    seg.set(p[i],1);
  }

  for(long long i=k;i<=n;i++){
    long long cur=(mod+bas-sub+add)%mod;
    res+=cur; res%=mod;
    if(i==n){break;}

    seg.set(p[i-k],0);
    sub+=(mod-seg.prod(0,p[i-k])); sub%=mod;
    sub+=seg.prod(p[i],n); sub%=mod;
    seg.set(p[i],1);
  }

  res*=modular_inverse(n-k+1); res%=mod;
  cout << res << "\n";
  return 0;
}