Atcoder python プログラミング競技プログラミング

累積和について~愚直に求めた場合と計算時間の比較~

久しぶりにAtcoderのコンテストに参加しました。Atcoder Beginner Contest 154
累積和の問題が解けなかったため、勉強のためにまとめました。

Atcoderの問題

今回自分が解けなかった問題はこちらです。

$N$ 個のサイコロが左から右に一列に並べてあります。左から $i$ 番目のサイコロは $1$ から $p_{i}$ までの $p_{i}$ 種類の目がそれぞれ等確率で出ます。

隣接する $K$ 個のサイコロを選んでそれぞれ独立に振ったとき、出る目の合計の期待値の最大値を求めてください。

こちらの問題の最大値を求めるところで累積和を使わないとTLEが起きてしまいます。
今回は自分が提出しTLEとなった解法と累積和を使った解法の計算時間を比較しました。

愚直な解法と累積和

今回提出してTLEとなった愚直な解法と累積和による解法を紹介します。

愚直な解法

まず初めに入力されたさいころの目を元にそれぞれの時の期待値を計算します。

さいころの目がpの時の期待値は$\frac{x+1}{2}$で求めることができます。

期待値を求めたうえで隣接するK個の期待値の和を算出しました。

例としてN=10、K=3の時の考えを示します。

オレンジのものが隣接する3つのサイコロを表していて、このオレンジのセルを順々に端まで移動します。この操作をするのにN-K+1回の処理が必要です。

また、この時オレンジのセルの和を求めます。和を求めるのにK回の処理が必要なので、

オレンジのセルを移動する処理、和を求める処理を合わせて(N-K+1）*K回の処理が必要になってしまいます。

pythonは$10^6$までの計算じゃないと間に合わないのですが、NとKの組み合わせ次第ではこれを超えてしまい間に合いません。

(補足)N = 200000の時の(N-K+1）*Kの最大値

蛇足感がありますが、N = 200000、Kを変数とした時に最大何回計算を行うかを求めました。

この式を式変形すると次のようになります。

$(\frac{200001}{2})^2 - (K - \frac{200001}{2})^2$

なので最大計算回数は$(\frac{200001}{2})^2$となり、愚直な方法ではpythonはおろかC言語でも間に合いません。

累積和を使った解法

累積和を使う場合は名前の通り、累積の和を求める前処理が必要です。

最初の配列には$p_0$を入れ、それ以降は前の項＋$p_n$を足していくという配列$S_n$を用意します。

この配列を使って任意の区間の和を求めます。

例えば左から4番目のセルから10番目のセルの間(下記図オレンジ区間)の和を求める場合は

$S_2 = p_0+p_1+p_2$

$S_9 = p_0+p_1+p_2+p_3+p_4+p_5+p_6+p_7+p_8+p_9$

なので、この二つの差は
$S_{9-2}= p_3+p_4+p_5+p_6+p_7+p_8+p_9$

となり、オレンジ色の区間の和を求めることができます。

この解法だと前処理に必要な配列$S_n$はO(N)で求められ、2つの差を求めるのにO(1)なので、愚直なやり方に比べ大分早く計算することができます。

ソースコード

今回はさいころの数nは問題文で与えられる最大値20000とし、K,pはrandom関数を使ってランダムな値を生成しました。試行回数は50回とし、50回試行した場合の平均計算時間を求めました。
ソースコードは下記の通りです。

import random
import time
import numpy as np
import seaborn
import matplotlib.pyplot as plt

#期待値を求める関数
#----------------------------------------------------------------------------------
def cal_ans(x):
    ans = (x+1)/2
    return ans
#----------------------------------------------------------------------------------

#pを作成する
#----------------------------------------------------------------------------------
def make_p_list(x):
    x = random.randint(1,1000)
    return x
#----------------------------------------------------------------------------------

time_list1 = []
time_list2 = []
k_list = []
ans_cs_list = []
ans_not_cs_list = []

for i in range(50):#50回トライする
    #入力
    #----------------------------------------------------------------------------------
    # n,k = map(int,input().split())#Atcoderで提出する際はこちら
    # p = list(map(int,input().split()))#Atcoderで提出する際はこちら
    n =  200000#Atcoderの問題で与えられている最大数
    k = random.randint(1,n)
    k_list.append(k)
    p = [0]*n
    p = list(map(make_p_list,p))#ランダムなpのリストを作成
    p = list(map(cal_ans,p))#入力されたものを元に期待値の計算
    #----------------------------------------------------------------------------------

    #累積和を使わない場合
    #----------------------------------------------------------------------------------
    start_time1 = time.time()
    ans_list = []
    ans = 0
    for i in range(n-k+1):
        if ans < sum(p[i:i+k]):#隣接するk個のさいころの期待値の和
            ans = sum(p[i:i+k])
    end_time1 = time.time()-start_time1
    time_list1.append(end_time1)
    ans_not_cs_list.append(ans)
    #----------------------------------------------------------------------------------

    #累積和を使った場合
    #----------------------------------------------------------------------------------
    start_time2 = time.time()
    S = [0]
    for i in range(n):
        S.append(S[i]+p[i])#累積和の計算
    ans = 0
    for i in range(n+1-k):#累積和を元に隣接するK個のさいころの期待値の和を求める
        if ans < S[i+k]-S[i]:
            ans = S[i+k]-S[i]#期待値の最大値の更新
    end_time2 = time.time()-start_time2
    time_list2.append(end_time2)
    ans_cs_list.append(ans)
    #----------------------------------------------------------------------------------

#結果の出力
#----------------------------------------------------------------------------------
print(time_list1)
print(time_list2)
print(sum(time_list1)/len(time_list1))#平均計算時間を求める
print(sum(time_list2)/len(time_list2))#平均計算時間を求める
print(k_list)#今回のkのリストを表示
print(ans_cs_list)
print(ans_not_cs_list)
print(ans_cs_list == ans_not_cs_list)#2つの解法で解が一致しているかの確認
plt.plot(range(len(time_list1)),time_list1,label = "NOT_CS")#累積和を使わなかった場合の結果
plt.plot(range(len(time_list2)),time_list2,label = "CS")#累積和を使った場合の結果
plt.grid()
plt.legend()
plt.show()
#----------------------------------------------------------------------------------

計算結果

計算時間は下記のようになりました。図の縦軸は時間[秒]、横軸は試行回数です。

また、青が愚直なやり方で実行した場合、オレンジが累積和を使った場合です。
(累積和は英語で"Cumulative sum"らしいので、今回は累積和のことを略してCSとしています。)

今回の例だと愚直なやり方だと平均60秒、累積和を使った解法だと平均0.06秒とおよそ100倍の違いが出ました。

愚直な方法でも1秒を切る場合がありましたが、最速でも0.5秒程度で10倍ほどの差があります。

今回の問題は実行時間は2 sまでなので、愚直なやり方ではTLEとなってしまいますが、累積和を使えばACすることができます。

参考文献

累積和を何も考えずに書けるようにする！

【累積和、しゃくとり法】初級者でも解るアルゴリズム図解

累積和について~愚直に求めた場合と計算時間の比較~

Atcoderの問題