AtCoder Beginner Contest 276 F,G 問題メモ

AtCoder Beginner Contest 276

F - Double Chance

• $N$ 枚のカードがあり、カード $i$ には整数 $A_i$ が書かれている
• $K=1,2,...,N$ について、次の問題を解く
• 問題
  • カード $1～K$ の中から、等確率でカードを1枚引き、戻すことを2回繰り返す
  • 1回目と2回目に引いた数をそれぞれ $x,y$ とすると、$\max(x,y)$ をその試行の得点とする
  • 得点の期待値を $\mod{998244353}$ で求めよ
• $1 \le N \le 2 \times 10^5$
• $1 \le A_i \le 2 \times 10^5$

たとえば $K=4$ で $1,4,7,10$ が書かれたカードから引く場合を考えると、 全試行パターンは1回目と2回目で4通りずつなので、$K^2=16$ 通り。

期待値は、全ての場合の得点を合計してこの $K^2$ で割ればよいので、合計の方を求められればよい。

得点ごとに、そのパターン数を考えると

得点が1になる場合:  1を2回連続で引く              1x1       = 1 通り
得点が4になる場合:  1,4から2回連続で引く場合から
                    得点が1になる場合を除く       2x2 - 1x1 = 3 通り
得点が7になる場合:  3枚から2回連続で引く場合から
                    得点が1,4になる場合を除く     3x3 - 2x2 = 5 通り
得点が10になる場合: 4枚から2回連続で引く場合から
                    得点が1,4,7になる場合を除く   4x4 - 3x3 = 7 通り

1x1 + 4x3 + 7x5 + 10x7 = 118

小さい順に $1,3,5,...$（奇数）通りずつパターン数が増えていく。得点とパターン数を掛け合わせればよい。
（同じ数字があった場合も、適当に順番を決めて扱えばよい）

で、この $K=4$ の値から $K=5$ の値を求められる。

$A_5=6$ を追加した場合、

  得点 パターン数
     1      1
     4      3
New! 6      5
     7   5→7
    10   7→9

3番目に追加されるのでパターン数は5となり、$6 \times 5$ が新たに増える。
また、それより大きい値の $7,10$ は、1つにつき $2$ ずつパターン数が増える。

なので、

• 追加する $A_i$ より小さい値が何個あるか $p$
• 追加する $A_i$ 以上となる値の合計はいくらか $q$

が、逐次求められればよいことになる。（$A_i \times (2p+1) + 2q$ だけ増える）

これは、$A_i$ をindexとして個数を管理するFenwick Treeと、合計を管理するFenwick Treeの2本を使って 1回あたり $\log{A_{max}}$ で求められる。

Python3

G - Count Sequences

• 以下の条件を満たす項数 $N$ の整数列 $A=(a_1,a_2,...,a_N)$ の個数を $\mod{998244353}$ で求めよ
• 条件
  • 要素は $0$ 以上 $M$ 以下で広義単調増加 $0 \le a_1 \le a_2 \le ... \le a_N \le M$
  • 隣り合う項（$i=1～N-1$ に対し、$a_i$ と $a_{i+1}$）は、3で割ったあまりが異なる
• $2 \le N \le 10^7$
• $1 \le M \le 10^7$

単調増加列の問題は、差分に着目すると上手くいくことがよくある。
末尾と $M$ との差分（下図 $b_5$）まで含めて考えると、差分の総和が上限値 $M$ となる。

A        0   a1   a2   a3   a4   M
Aの差分   b1 + b2 + b3 + b4 + b5    = M

隣り合う項の3で割ったあまりが異なるので、両端を除く $b_2～b_4$ は、3の倍数であってはいけないということになる。

逆にこれさえ満たせば条件を満たす数列となるため、そのような非負整数列 $B=(b_1,...,b_{N+1})$ の個数を数えればよい。

ここで、$B$ を3で割った商 $X=(x_1,...,x_{N+1})$ とあまり $Y=(y_1,...,y_{N+1})$ に分けて考える。

条件は「$y_2～y_N$ は $0$ 禁止（$1,2$ のいずれか）」「$3x_i+y_i$ の総和が $M$」となる。

条件を満たす $X,Y$ は、例えば以下のようにして構築できる。

• $y_1$ を $0,1,2$ の中から決める
• $y_2～y_N$ は最低でも1なのであらかじめ1を敷いておき、「2になる要素がどれか」を決める
• ここまでに決めた数によって、$M$ の残りから、$X$ の総和と $y_{N+1}$ は自動的に決まる
• $X$ の総和の $x_1,...,x_{N+1}$ への割り振り方を決める

以下の2つを固定すると、

• $s$: $y_1$ の値（$0/1/2$ の3通り）
• $k$: $y_2～y_N$ のうち、$2$ になる要素の個数（$0～N-1$ の $N$ 通り）

そのパターン数が $O(1)$ で求められる。

• $y_2～y_N$ のうちの $2$ の配置の決め方は ${}_{N-1}C_{k}$
• $X$ の総和は $\dfrac{M-(N-1)-s-k}{3}$（切り捨て）に決まる。これを $t$ とする
  • $x_1,...,x_{N+1}$ への割り振り方は ${}_{N+1}H_{t}$

この2つを掛け合わせたものが、$s,k$ を固定した時のパターン数となる。全探索して合計すれば答え。

本番では3で割った商とあまりに分割するところまでは思いついたが、$b_1$ を特別視しすぎて $b_2～b_N$ だけで $X,Y$ を考え、$b_1$ は別で計算してしまったため、考慮すべきパターン数が増え、計算量が $NM$ から減らせなかった。

Python3