Codeforces Round 829 (Div. 2) - E. Wish I Knew How to Sort (2000)

E. Wish I Knew How to Sort

$0$이 $g$개 있다고 하자.

$dp[k]$ 를 처음 $g$개의 위치에 $k$ 개의 $0$이 있기위한 기댓값이라고 하자.

$dp[g]=0$ 이다. 모든 $0$이 앞에 $g$ 개에 모두 있다면 sorted 상태이기 대문이다.

점화식을 다음과 같이 세울 수 있다.

$dp[k]$가 있을 때, 앞에 $g$ 개에서 $1$ 인것을 고르고 나머지 $(n-g)$ 에서 $0$ 인것을 골라야 상태가 전이된다.

이 때의 확률은 $p_k=\dfrac{(g-k) \cdot (g-k)}{\binom{n}{2}}=\dfrac{2(g-k)^2}{n \cdot (n-1)}$ 이다.

앞에 $g$ 개에서 $1$ 은 현재 $g-k$ 개가 있고, 결국 뒤에서도 앞으로 와야할 $0$은 $(g-k)$ 개가 있을 것이기 때문이다.

따라서 $dp[k]=p_{k} \cdot dp[k+1]+(1-p_k) \cdot dp[k]+1$ 이다.

• $p_k$의 확률로 앞 $g$개에 $0$가 $k+1$개 있는 상태로 넘어갈 수 있으므로 그 때의 기댓값과 $p_k$ 를 곱한 값을 더한다.
• $(1-p_k)$ 의 확률로 시행이 무의미하게 써지기 때문에 $(1-p_k) \cdot dp[k]$ 를 더한다.
• 우린 이번으로 한 번의 시행을 했기 때문에 $1$ 을 더한다.

이걸 풀면

$dp[k] \cdot p_k=p_k \cdot dp[k+1]+1$

$dp[k]=dp[k+1]+\dfrac{1}{p_k}=dp[k+1]+\dfrac{n(n-1)}{2(g-k)^2}$ 이다.

처음에 $g$ 개의 위치에 $0$ 이 있는 개수를 $x$ 라고 하면 $dp[g]=0$ 이고 정답은 $dp[x]$ 이다.

const ll mod = 998244353;  
inline ll md(ll x) { return md(mod, x); }  
inline ll poww(ll a, ll b) {  
   ll ret = 1;  
   while (b) {  
      if (b & 1) {  
         ret *= a;  
         if (~mod) ret %= mod;  
      }  
      a *= a;  
      if (~mod) a %= mod;  
      b >>= 1;  
   }  
   return ret;  
}  
inline int inv(int x) {  
   return poww(x, mod - 2);  
}  
void solve() {  
   int n;  
   cin >> n;  
   vi a(n);  
   fv(a);  
   int g = cntt(a, 0), x = 0;  
   for (int i = 0; i < g; i++) if (a[i] == 0) x++;  
  
   if (x == g) {  
      cout << 0 << endl;  
      return;  
   }  
  
   int val = 0;  
   for (int k = g - 1; k >= x; k--) {  
      val = md(val + md(md(md(md(n * (n - 1)) * inv(2)) * inv(g - k)) * inv(g - k)));  
   }  
   cout << val << endl;  
}