この記事は拙著「AtCoder 凡人が『緑』になるための精選50問詳細解説」のサンプルです
価格：100円
kindle：https://www.amazon.co.jp/dp/B09C3TPQYV
booth(pdf)：https://sano192.booth.pm/items/3179185
前：https://qiita.com/sano192/items/261790cbd90dbe6cf596
次：https://qiita.com/sano192/items/1b0c9f6a0794e5595ef0
【目標】
・期待値の計算ができるようになる
・数列の区間和を累積和を使って高速で計算できるようになる
【概要】
期待値、区間和の知識を求められる問題。これらは非常によく出題されるテーマのため、逆に考えるとこの問題は2つの知識を同時に身につけるチャンスでもある。難しいと感じるかもしれないが、じっくり解説を読めばそこまでややこしい問題ではないので、頑張っていこう。
【方針】
各サイコロについて出る目の期待値を計算し、K個ずつ和を計算して一番大きいものを出せば良い。
最初に期待値について説明しよう。
期待値はよく「出ることが期待される数」などと説明されるが厳密には少し意味が違う。しかし厳密に説明するとやたら難しい数学の話になってしまう。
本問を解くにあたっては「サイコロの目の期待値は(目)*(確率)を全部足す」程度の雑な理解でよい。きちんと理解したい人は「期待値」で検索しよう。以下のページでは数学的に厳密な説明がされている。
https://physnotes.jp/stat/drv_crv/
具体的な計算方法を説明する。
サイコロの目の期待値は(目)*(確率)を全パターン足せば計算できる。
6面なら
目：1 * 確率：1/6 = 1/6
目：2 * 確率：1/6 = 1/3
目：3 * 確率：1/6 = 1/2
目：4 * 確率：1/6 = 2/3
目：5 * 確率：1/6 = 5/6
目：6 * 確率：1/6 = 1
これらを全て足すと
1/6+1/3+1/2+2/3+5/6+1=7/2=3.5
となる。
サイコロが複数ある場合、出る目の和の期待値は単純に足し算で求められる。
つまりサイコロ2つの目の和の期待値は
3.5+3.5=7
となるし、サイコロ3つの目の和の期待値は
3.5+3.5+3.5=10.5
となる。
サイコロがp面の場合、期待値は
目：1 * 確率：1/p = 1×(1/p)
目：2 * 確率：1/p = 2×(1/p)
目：3 * 確率：1/p = 3×(1/p)
...
目：p * 確率：1/p = p×(1/p)
これらを全て足して
1*(1/p)+2*(1/p)+3*(1/p)+...+p*(1/p)=(1+2+3+...+p)(1/p)
ここで(1+2+3+...+p)は和の公式を使うと(1/2)p(1+p)となるから期待値は
(1/2)p(1+p)(1/p)=(1+p)/2
となる。
K個のサイコロについて目の合計の期待値は、それぞれのサイコロの期待値を足せばよいだけ。
入力例1を使ってどのように計算を行うか具体的に説明する。
「入力例1」
5 3
1 2 2 4 5
N=5(サイコロの数)
K=3(区間)
p1=1
p2=2
p3=2
p4=4
p5=5
まずp1、p2、p3、p4、p5それぞれについて出る目の期待値を計算する。
p面サイコロの出る目の期待値は(1+p)/2だから
p1：p=1→(1+1)/2=1
p2：p=2→(1+2)/2=1.5
p3：p=2→(1+2)/2=1.5
p4：p=4→(1+4)/2=2.5
p5：p=5→(1+5)/2=3
K=3なので、隣接する3つのサイコロについて目の期待値を足し算する。
p1~p3=1+1.5+1.5=4
p2~p4=1.5+1.5+2.5=5.5
p3~p5=1.5+2.5+3=7
故に期待値は3~5番目のサイコロを振ったときが最大で答えは7。
期待値は計算できるようになったがまだこの問題は解けない。
K個の隣接するサイコロについて期待値を愚直に足し算していくとTLE(実行時間制限超過)するからだ。
そこで区間和の高速計算が必要となる。
区間和とは「ある数列に対して一部の区間の和」を意味する。
例えば
1 2 4 5 7 4 2
という数列について、最初の1を「0番目」としたとき、「2番目」～「５番目」の区間和は
4+5+7+4=20
と計算できる。
区間和についてはあらかじめ累積和を計算しておくことで高速に計算できるテクニックがある。
まず累積和を計算する。「X番目」までの累積和とは「0番目」～「X番目」までの区間和を意味する。
1 2 4 5 7 4 2について
「0番目」までの累積和=「0番目」～「0番目」の区間和=1
「1番目」までの累積和=「0番目」～「1番目」の区間和=1+2=3
「2番目」までの累積和=「0番目」～「2番目」の区間和=3+4=7(1+2+4)
「3番目」までの累積和=「0番目」～「3番目」の区間和=7+5=12(1+2+4+5)
「4番目」までの累積和=「0番目」～「4番目」の区間和=12+7=19(1+2+4+5+7)
「5番目」までの累積和=「0番目」～「5番目」の区間和=19+4=23(1+2+4+5+7+4)
「6番目」までの累積和=「0番目」～「6番目」の区間和=23+2=25(1+2+4+5+7+4+2)
例えば「0番目」～「3番目」までの累積和(1+2+4)を計算したい時、
「2番目」までの累積和+「3番目」
を計算すればよいから
3+4=7
と計算している。
累積和を使って区間和を計算しよう。
例えば「2番目」～「５番目」の区間和(4+5+7+4)は
「2番目」～「５番目」の区間和
=(「5番目」までの累積和)-(「1番目」までの累積和)
=23-3
=20
と計算できる。
実際に数値を入れて確認しよう。
「2番目」～「５番目」の区間和
=(4+5+7+4)
=(「5番目」までの累積和)-(「1番目」までの累積和)
=(1+2+4+5+7+4)-(1+2)
と、このように不要な部分をきれいに引き算できることがわかる。
一般に「X番目」～「X+K番目」の区間和は
・X=0の場合
「X+K番目」までの累積和
・1<=Xの場合
(「X+K番目」までの累積和)-(「X-1番目」までの累積和)
と計算できる。
累積和の計算はO(N)。
累積和を計算した後の区間和の計算はO(1)。
このテクニックを使えばTLE(実行時間制限超過)しない。
【実装】
入力を受け取る。
N,K=map(int, input().split())
p=list(map(int, input().split()))
それぞれのサイコロについて、期待値を計算したリスト(p_ex)を作る。
期待値はp面のサイコロについて(1+p)/2と計算できる。
p_ex=[]
for i in range(N):
p_ex.append((p[i]+1)/2)
次にp_exの累積和を計算する。
累積和を計算するときはまず0番目の数をリストに入れ、i=1,2,...,N-1について
「i番目までの累積和」=「i-1番目」までの累積和+「i番目」
を順に計算していく。
名前はp_ex_Cumとする。(累積和は英語でCumulative sum)
p_ex_Cum=[p_ex[0]]
for i in range(1,N):
p_ex_Cum.append(p_ex_Cum[i-1]+p_ex[i])
累積和が計算できたらK個のサイコロについて区間和を計算し、より大きいものを答えとして更新していく。
・「0番目」～「0+K-1番目」の区間和=「0+K-1番目」までの累積和
・「i番目」～「i+K-1番目」の区間和=(「i+K-1番目」までの累積和)-(「i-1番目」までの累積和)
と計算できる。
ans=-10**15
for i in range(N-K+1):
if i==0:
ans_tmp=p_ex_Cum[i+K-1]
else:
ans_tmp=p_ex_Cum[i+K-1]-p_ex_Cum[i-1]
ans=max(ans,ans_tmp)
ans_tmpはanswer temporary(一時的な答え)の意味。
最後に答えを出力して終わり。
print(ans)
【コード全文】
N,K=map(int, input().split())
p=list(map(int, input().split()))
p_ex=[]
for i in range(N):
p_ex.append((p[i]+1)/2)
p_ex_Cum=[p_ex[0]]
for i in range(1,N):
p_ex_Cum.append(p_ex_Cum[i-1]+p_ex[i])
ans=-10**15
for i in range(N-K+1):
if i==0:
ans_tmp=p_ex_Cum[i+K-1]
else:
ans_tmp=p_ex_Cum[i+K-1]-p_ex_Cum[i-1]
ans=max(ans,ans_tmp)
print(ans)
この記事は拙著「AtCoder 凡人が『緑』になるための精選50問詳細解説」のサンプルです
価格：100円
kindle：https://www.amazon.co.jp/dp/B09C3TPQYV
booth(pdf)：https://sano192.booth.pm/items/3179185
前：https://qiita.com/sano192/items/261790cbd90dbe6cf596
次：https://qiita.com/sano192/items/1b0c9f6a0794e5595ef0