At Coder 東京海上プログラミングコンテスト　C問題　lamps

前提・実現したいこと

C問題をTLEなく正解したい

制約
1≦N≦2×10＾5
1≦K≦2×10＾5
0≦Ai≦N

発生している問題・エラーメッセージ

TLE
（計算量は４２×KでTLEにはならないのでは？と思うのですが、、、）

該当のソースコード

python
1from itertools import accumulate
2import numpy as np
3
4n, k = map(int, input().split())
5a = list(map(int, input().split()))
6
7for q in range(min(42,k)):
8  b = [0]*(n+1)
9  for i in range(n):
10    l = max(0, i-a[i])
11    r = min(i+a[i]+1, n)
12    b[l] += 1
13    if r <= n-1: b[r] -= 1
14  
15  a = list(accumulate(b))
16  a = a[:-1]
17   
18  if np.all(a == n):
19    break
20
21print(*a)
22

試したこと

・いもす法の活用
・ループ数の短縮（最大４２回以下）

行動規範の内容に同意します

回答1件

ベストアンサー

pythonは数値計算がとても遅い言語です。
何故かというと静的な型を持たないので、計算が発生するたびに双方の数値（数値じゃないかもしれない）の型を調べ、型がそろっていないなら適切にキャストまたは変換し、型に対して適切な演算関数を探して…という
型チェック処理にとても時間がかかるからです。このコードだとmaxとminが特に遅いです。

この型チェック地獄を解決するアプローチが主に2つあります。

numpyの利用

numpyで配列を作成する際にdtypeで型を指定します。するとnumpy配列同士の計算においては型チェックを行う必要がなくなるのでとても高速になります。またnumpy内の計算処理がC言語で書かれているためその点でも計算が高速です。
numpyで高速化するには、実装したい計算をうまく行列計算できるように組み替える必要があります。

numba.njitの利用

numbaはpythonのコードをネイティブコードに変換して実行するライブラリです。njitを利用すると変数の型を固定してを高速に演算できます。AtCoderでも使えます。
コンパイルエラーなく利用するにはそれなりに流儀に従ったコードを書く必要があります。

pythonでACを取れている方はほとんどがどちらかの方法で解決しています。njitを使っている方のほうが多そうですね。

質問のコードを極力そのままにnjit対応コードに修正するとこんな感じになるかと思います。
最悪パターンでも1秒かかりません(AtCoderコードテストにて)。

python
1%%time
2from itertools import accumulate
3import numpy as np
4from numba import njit
5
6n, k = map(int, input().split())
7a = np.array(list(map(int, input().split())),dtype=np.int64)
8# 大規模データテスト用
9#n = 2 * (10 ** 5)
10#k = 2 * (10 ** 5)
11#a = np.zeros(n,dtype=np.int64)
12
13#時間のかかる関数をjit化
14@njit
15def loop1(a):
16    #b = [0]*(n+1)
17    b = np.zeros(n+1, dtype=np.int64)
18    for i in range(n):
19        l = max(0, i-a[i])
20        r = min(i+a[i]+1, n)
21        b[l] += 1
22        if r <= n-1:
23            b[r] -= 1
24    b = np.cumsum(b)[:-1]
25    return b
26
27for q in range(min(42,k)):
28    a = loop1(a)
29
30# outが大変なことになるので普段はコメントアウト
31#print(*a)