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

N個のサイコロが左から右に一列に並べてあります。左から 

i番目のサイコロは 1から $p_i$までの $p_i$種類の目がそれぞれ等確率で出ます。

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

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

累積和を使った解法

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()
#----------------------------------------------------------------------------------