ARC102 E - Stop. Otherwise...(700)

ARC102 E - Stop. Otherwise..(700)

E - Stop. Otherwise...
コンテスト中にすごい解けそうな気がしたんだけどなぁ…
多分想定解ではないです．
(editionalも解説放送もあっさりしてて正直想定解私には理解できなかった)

問題概要

K面サイコロをN回振ったときに各i(2<=i<=2K)に対し，
・どの異なる2つのサイコロの出目の和もiにならないような出目の組の場合の数
を998244353でわった余りを求めよ．

制約

• 1≤K≤2000
• 2≤N≤2000
• K,Nは整数である

考えたこと

とりあえずサンプルから見る感じ
・i<=Kの範囲で考えてひっくり返してくっつければ良さそう．
・i%2==0の時，iとi+1は一致するっぽいので偶数のiだけ考えれば良さそう．
とりあえずそういう†前提†(証明はない)のもと考えていく

i=2の時

まずとりあえずi=2の時で考える． 区別のないサイコロを振るとかちょっと難しいので1～Kまでの値のついた箱にボールをN個入れる入れ方の場合の数と言い換えて考える．

またこの時どの異なる2つのサイコロの出目の和もiにならないという条件は， 任意の球を二つ取り出した時にその箱の和がiとなる組み合わせが存在しないと言い換えられる．
今回はi=2なので出目の和が2となるのは1の箱から2つ取り出した時だけなので，1の箱に入っている球の数が0個か1個なら良い．

1の箱に球がx個入ってる時，
残りの(K-1)個の箱の中に球をN-x個入れる入れ方なので

と書ける(球全部と箱-1個の線の並べ方みたいなイメージ(中学数学でやった気がする))
なのでx=0とx=1の時の値を足したものがi=2の時の場合の数

i >2の時

最初の考察よりi%2==0としてもよい．
考えやすいように具体的にN=5,K=6,i=6と値を与える．

和がiとなる組が存在しないためには， iは偶数なのでi//2の箱には1個以下，また1<=x<i//2の全てのxに関して{x,i-x}の組の少なくともどちらかが0個となっている必要がある．(今回は{1,5}と{2,4}の組，(図で線で結んであるもの)
これを数え上げればよいのだが，少なくとも一方が0個という考え方だとどちらも0個となるものを重複して数えてたりしてしまうためヤバミが深い(包除原理上手いこと使えばできるんだろうか…)
ので，とりあえずどちらも0個となっている組を考えず，この{1,5},{2,4}の組み合わせの内一方が0個であればもう一方が1個以上と仮定する． またどちらも0個の組というのを便宜上0pairと書く．

0pairが0組

{1,5},{2,4}の組のうち0となるもの(どちらか一つ)を選ぶ方法は2²=4通りである．

例えば1,2の箱の中身が0個だったとすると，4,5の箱中身は1個以上入っている必要がある．これを先にそれぞれに突っ込んでおけば図のようになるが，
・1,2の箱は使わないためK=4
・球を既に2ついれたためN=3
・3の箱に2つ以上入れてはいけないという条件なのでi=2と同じ．
とそれぞれ置き換えることができる． 先ほどi=2となるものについての考察よりこの場合の数はわかり，また0となるものに依らず同じことがいえるので*=4をしたものが，この時の合計の場合の数になる．

0pairが1組

{1,5},{2,4}の組のうち，どちらも0個となる組が一つあるとする．どちらも0pairの選び方は自明に通りある．これを消すと，

となり，図のようにN=5,K=4,i=4と同様な状況となる． ここで重要なのは，0pairは全体で1個なので他の組(例えば今回なら{2,4})には0pairは存在しえない． つまりN=5,K=4,i=4で 0pairが0組なので，この場合の数は前に述べた方法で求められる． また0pairの選び方に依らず同様の考え方で求められるので*=2したものがこの時の合計の場合の数となる．

0pairの数を増やしていっても同様に求めることができる．
また0pairの数が固定であるためこれらの場合の数で重複しているものはなく，これらを足し合わせたものがこの時の場合の数となる． これを実装します．(辛そう

実装

とりあえず上で考察したことを式に書いてまとめる．

：球の数n個，箱の数k個，和がiにならないように取る場合の数
：のうち0pairの数がp個となる場合の数
こんな感じでfとgを定義

• fについて

  • if i == 2 :
    
  • if i > 2 :
    

• gについて

  • if p == 0 :
    
  • if p > 0 :
    

ここでは片方が0個でなければならない箱の組の数()を表す

このように全てi=2のケースに帰着できるので数え上げることができる．
ただし例外処理としてNが少なくKが大きい時，この計算途中でnが負となるケースが出てくるので，

• if n == 0 : (全ての箱に球を入れない一通り)
• if n < 0 : （そのような場合の数は存在しえない）
• if k==1 and n>1 and i==2 : (1面サイコロを2回以上振ったら和が必ず2になる)

とする必要がある．(私はこの例外処理でミスって本番中に通せませんでした…

あとは上の式通りにfとgを実装して2<=i<=2*Kの範囲でまわすだけです．
(ここで一番最初に†前提†として述べた前半と後半が対称性により一致するのと，i%2==0の時i,i+1が一致するのを使う) あとまあmodとクソデカCombinationくんが出てくるので，逆元使って事前計算してO(1)にしておきましょうって感じです(これはAtCoderでも過去に何回も出てる話だし有名なものなので『Combination 高速 逆元 [検索]』とかで検索すればこんな日記よりよっぽどわかりやすい記事が出てくると思うので省略します)
今回最大N+Kなのでそのへんまでテーブル作っておけばよい
Submission #3128053 - AtCoder Regular Contest 102

K,N = list(map(int, input().split()))
MAX = K+N+3
mod = 998244353

fac = [1]*(MAX+1)
for i in range(1,MAX+1):
    fac[i] = (fac[i-1]*i)%mod

rev_m = [1]*(MAX+1)
rev_m[MAX] = pow(fac[MAX],mod-2,mod)
for i in range(MAX,0,-1):
    rev_m[i-1] = (rev_m[i]*i)%mod

def Comb(n,k):#nCk
    return (fac[n]*rev_m[k]*rev_m[n-k])%mod


def f(n,k,i):
    if i == 2:
        if n < 0 or (k==1 and n>=2): #例外処理
            return 0
        elif n == 0:
            return 1
        else:
            return (Comb(k-2+n,n)+Comb(k-2+n-1,n-1)) % mod
    else: #i>2
        Sp = i//2-1
        ans = 0
        for p in range(Sp+1):
            ans += g(n,k,i,p)
        return ans % mod

def g(n,k,i,p):
    Sp = i//2-1
    if p == 0:
        return f(n-Sp,k-Sp,2) * pow(2,Sp,mod) % mod
    else: #p>0
        return g(n,k-2*p,i-2*p,0) * Comb(Sp,p)

ans = []
for i in range(2,K+2):
    if i%2 == 0:
        tmp = f(N,K,i)
    ans.append(tmp)
    print(tmp)

ans = ans[::-1]
for i in range(1,K):
    print(ans[i])

感想

わざわざ別記事にしておいてこんな謎解法置いておくの申し訳のなさ感ある．
まあ一応これでACは取れるし，通れば正義なので(?)こんな考え方もありますよってことで
個人的には頑張って考察したら最終的にまあまあ綺麗な感じで書けたのでわりと達成感の感じた解法です(は? (†前提†が成り立つ理由考察してませんが…