pythonは数値計算がとても遅い言語です。
何故かというと静的な型を持たないので、計算が発生するたびに双方の数値(数値じゃないかもしれない)の型を調べ、型がそろっていないなら適切にキャストまたは変換し、型に対して適切な演算関数を探して…という
型チェック処理にとても時間がかかるからです。このコードだとmax
とmin
が特に遅いです。
この型チェック地獄を解決するアプローチが主に2つあります。
numpyで配列を作成する際にdtype
で型を指定します。するとnumpy配列同士の計算においては型チェックを行う必要がなくなるのでとても高速になります。またnumpy内の計算処理がC言語で書かれているためその点でも計算が高速です。
numpyで高速化するには、実装したい計算をうまく行列計算できるように組み替える必要があります。
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)
余談
どちらの方法を使わなくてもACできてる方もいる。ギリギリだけど。
バッドをするには、ログインかつ
こちらの条件を満たす必要があります。
2020/06/21 01:07