(MojaCoder) Mex Appends 解法メモ

面白かったので
mojacoder.app

問題概要

0 以上 $K-1$ 以下の要素からなる長さ $N$ の数列 $A$ がある。はじめ空の集合 $S$ に対して、
- (1) $A_i$ を $S$ に追加。
- (2) $mex(S)$ を答えに加算。
を $N$ 回行う。考えうる全ての数列 $A$ に対して答えを求めて総和を出力せよ。

$1\le N\le 9\times 10^{8}$
$1\le K\le 10^{5}$

考察

寄与を求めていく系っぽいのでまず $1\le x\le K$ が mex になる回数を数えることにする。
$i\ (1\le i\le N)$ 回目の操作において mex = $x$ である条件は、以下のように整理できる:

1 ~ $i-1$ 回目に追加された要素に、 $1,\ 2,\ \cdots ,x-1$ が少なくとも１つ以上含まれ、かつ $x$ は含まれていない

すると包除原理が見えるので、とりあえず以下のような式が立った:

$x$ が $i$ 回目に追加される回数に重み $x$ をかけた寄与は、 $\displaystyle \sum_{m=0}^{x-1} (-1)^{m} \times \binom{x-1}{m}\times (K-1-m)^{i} \times K^{N-i}\times x$

最後の $K^{N-i}$ は $i+1$ 番目以降は何でも良いのでこれが付く（解いた時は割としばらく忘れていた...）

しかし、何やら $x-1$ とか $K-1-m$ とか例外処理が必要そう + $x$ が式にたくさん登場しているので和をとるのが面倒。

これをさらに考えると、「ちょうど $x$ 」ではなく、「mex が $x$ 以上」とすると何やらうまくいく。
こうすると、mex = $x$ になるパターンは「 $1$ 以上、 $2$ 以上、...、 $x$ 以上」のちょうど $x$ 回重複して数えられるので、重みをつけなくて良いことになる。（writerさんの解説では割とあっさり書いてたけど天才ポイントな気がする。実は典型？）

すると式は少し良い感じになって、
- $\displaystyle \sum_{m=0}^{x} (-1)^{m} \times \binom{x}{m}\times (K-m)^{i} \times K^{N-i}$

となる。

このままだと $O(NK^{2})$ なので、あとは $\sum$ を解消すれば良い。~~じゃあ WolframAlpha さんよろしく。~~

式変形

どうやらこのままでは解消するのが難しいようなので、工夫を考える。改めて式を書くと、

$\displaystyle \sum _ {x=1}^{K} \sum _ {i=1}^{N} \sum _ {m=0}^{x} (-1)^{m} \times \binom{x}{m}\times (K-m)^{i} \times K^{N-i}$

まず $i$ については順番を変えて普通に計算できて、 $m\gt 0$ では

$\displaystyle \sum _ {x=1}^{K} \sum _ {m=1}^{x} (-1)^{m} \times \binom{x}{m}\times \dfrac{(K-m)}{m}\times (K^{N}-(K-m)^{N})$

となる。（ $m=0$ は簡単なので、別で計算する） $m$ は式が複雑で解消しづらいようだが、ここで形式的に $m \gt x$ を考えてみる。
$n\lt r$ では $\binom{n}{r}=0$ であるので、答えの値は変わらないが、これによって

$\displaystyle \sum _ {x=1}^{K} \sum _ {m=1}^{K} (-1)^{m} \times \binom{x}{m}\times \dfrac{(K-m)}{m}\times (K^{N}-(K-m)^{N})$

となって、２番めの $\sum$ は $x$ に依存しなくなった。よって順番は交換できて、 $x$ についてならすぐ解消できる。

$\displaystyle \sum _ {m=1}^{K} (-1)^{m} \times \binom{K+1}{m+1}\times \dfrac{(K-m)}{m}\times (K^{N}-(K-m)^{N})$

これで $O(K\log N)$ で計算できる形になった。

感想

よく見る $\displaystyle \sum _ {i=1}^{N-1} \sum _ {j=i+1}^{N} f(i,\ j)$ みたいなのを $\displaystyle \sum _ {i=1}^{N} \sum _ {j=1}^{N} f(i,\ j)$ にするテクの仲間っぽさを感じる。
$\binom{n}{r}=0\ (n\lt r)$ だから足して良いというのはあまり見たことが無かったので面白かった。