Python のタイムアウトについて

Question

HackerRankというサイトでPythonの問題を解いていたのですがそこでPythonの実行時間が長すぎるため停止しましたと言われました。
以下が問題になります。（英語です）



これに対して私が出したコードはこちらになります。私が書いたコードはdefのarrayManipulationの部分のみです。ここが原因で時間が足りないと思われます。また実行時間が長すぎると言われたテストケースはこちらになります。どうすれば改善できるでしょうか？
テストケース

#!/bin/python3
import math
import os
import random
import re
import sys

# Complete the arrayManipulation function below.
def arrayManipulation(n, queries):
    list0=[0 for i in range(n)]
    for query in queries:
        bet1=[j-1 for j in query[0:2]]
        for i in range(bet1[0],bet1[1]+1):
            list0[i]=list0[i]+query[2]
    return max(list0)

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')
    nm = input().split()
    n = int(nm[0])
    m = int(nm[1])
    queries = []
    for _ in range(m):
        queries.append(list(map(int, input().rstrip().split())))
    result = arrayManipulation(n, queries)
    fptr.write(str(result) + '\n')
    fptr.close()

Comments

[8524ba23]

問題は画像ではなくテキストとして提示すると回答が得られやすくなるかと思います。

Answer 1

※下記は
https://sites.northwestern.edu/acids/2018/11/12/solution-hackerrank-array-manipulation/
について筆者なりの解説を加えたものになります。

質問者さんのコード自体は、正しい回答を導き出せるコードにはなってるのですが、時間がかかりすぎるということのようですね。

        bet1=[j-1 for j in query[0:2]]
        for i in range(bet1[0],bet1[1]+1):
            list0[i]=list0[i]+query[2]

で「与えられた範囲」の「すべて」の要素に同じ数を足しているところが特にボトルネックとなっているのではないでしょうか。

足しこむのは「同じ数」なので、範囲の「左右端っこ」さえデータ反映すればいいですね。

たとえば
時間切れになった問題の最初のクエリ

5216499 7952900 40114

について、7952900 - 5216499 + 1 =2736402 個の要素全部に対して足し算を行うよりも
「5216499」番目と「7952900」番目の「２つ」だけについて、データを反映させれば、2736402/2≒136万分の1以下の処理量で済みそうです。

ここで、各要素の変化量に着目します。

足し算を行った後の配列を見ると、範囲内に対して同じ数を足しているので、範囲内の変化量はゼロですが、範囲の端っこでは変化量が生じています。

（上で言えば
1番目から5216498番目までの変化量はゼロ、
5216498番目から5216499番目の変化量は　+40114、
5216499番目から7952900番目までの変化量はゼロ
7952900番目から7952901番目の変化量は、-40114
7952902番目から10000000番目までの変化量はゼロ
となっています）

よって、左隣の要素に対する変化量のみを記録するようにしてみましょう。

例えば、10個の要素の配列
0, 0, 0, 0, 0, 0, 0, 0, 0, 0
に対して、
クエリ2 5 100
が与えられたときは、そのまま計算すると
0, 100, 100, 100, 100, 0, 0, 0, 0, 0
となりますが、
それぞれの左の数に対する変化量は
0, 100, 0, 0, 0, -100, 0, 0, 0, 0
となっています。

続けて、
クエリ3 4 50
が与えられたとすると、そのままの計算では
0, 100, 150, 150, 100, 0, 0, 0, 0, 0
となりますが、
各要素の左に対する変化量は、
0, 100, 50, 0, -50, -100, 0, 0, 0, 0
となっています。

一般化すると、

クエリquery: (i1, i2 , value)について

list0[i1-1] += value
list0[i2] -= value

とすればよいことになります。

最終的に最大値を求めるときには、
左から順番に足し算を行い、足し算の結果をそれまでの最大値と比較し、
大きければそれを最大値とし、
最後まで残った最大値が答えとなります。

以上を反映したコードが以下になります。
コード出典：https://sites.northwestern.edu/acids/2018/11/12/solution-hackerrank-array-manipulation/

def arrayManipulation(n, queries):
    arr = [0]*n
    for i in queries:
        arr[i[0] - 1] += i[2]
        if i[1] != len(arr): #  与えられたクエリの右端が配列の右端と重なる場合をトラップ
            arr[i[1]] -= i[2]
    maxval = 0
    itt = 0
    for q in arr:
        itt += q
        if itt > maxval:
            maxval = itt
    return maxval

Comments

[Charowan]

非常にわかりやすい解説ありがとうございます！！

Answer 2

        for i in range(bet1[0],bet1[1]+1):
            list0[i]=list0[i]+query[2]

どの程度の時間制限があるのかわかりませんが、テストケースのパラメーターを見る限り、ここのループはなくす必要があるでしょう。
方法としてはimos法とかBinary Indexed TreeとかSegment Treeとかクエリの端点をソートして足していくとかが考えられます。