東京海上日動プログラミングコンテスト2020 C問題解けず…

東京海上日動プログラミングコンテスト2020、２完して少しレートが上がりました。３問目は、プログラム自体は書けたのですが、タイムアウトしてしまいました。終わった後も検討したのですが、改善せず。誰か教えてください！！

f:id:poohandduffy:20200614114052p:plain

問題文

数直線上に電球が $N$ 個並んでおり、電球には左から順に $1$ から $N$ までの番号がついています。電球 $i$ は座標 $i$ にあります。

電球には光の強さを表す非負整数値が定まっており、座標 $x$ に光の強さ $d$ の電球があるとき、その電球は座標 $x - d - 0.5$ から座標 $x + d + 0.5$ までの区間を照らします。初めは電球 $i$ の光の強さは $A_{i}$ です。そこで、以下の操作を $K$ 回繰り返し行います。

$1$ 以上 $N$ 以下の各整数 $i$ に対し、操作時に座標 $i$ を照らしている電球の個数を $B_{i}$ とする。そして、各電球 $i$ の光の強さを $B_{i}$ に変更する。

$K$ 回の操作を行った後の各電球の光の強さを求めてください。

制約

$1 ≦ N ≦ 2 \times 10^{5}$
$1 ≦ K ≦ 2 \times 10^{5}$
$0 ≦ A_{i} ≦ N$

入力

入力は以下の形式で標準入力から与えられる。

 $N$   $K$ 
 $A_{1}$   $A_{2}$   $\dots$   $A_{N}$

出力

$K$ 回の操作を行った後の電球 $i$ の光の強さ $A {^{'}}_{i}$ を、以下の形式で標準出力に出力せよ。

 $A {^{'}}_{1}$   $A {^{'}}_{2}$   $\dots$   $A {^{'}}_{N}$

入力例 1 Copy

Copy

5 1
1 0 0 1 0

出力例 1 Copy

Copy

1 2 2 1 2

始めに座標 $1$ を照らしている電球は電球 $1$ のみであるので、操作後の電球 $1$ の強さは $1$ になります。また、始めに座標 $2$ を照らしている電球は電球 $1$ と電球 $2$ であるので、操作後の電球 $2$ の強さは $2$ になります。

入力例 2 Copy

Copy

5 2
1 0 0 1 0

出力例 2 Copy

Copy

3 3 4 4 3

イモス法で解いたのですが…


import sys
n, k = map(int, sys.stdin.readline().rstrip().split())
an = list(map(int, sys.stdin.readline().rstrip().split()))

import numpy as np

t= min(k, 45)

for _ in range(t):
　bn = [0]*n
　for i in range(n):
　　bn[max(i - an[i], 0)] += 1

　　if i+ an[i]+1 <n:
　　　bn[i+an[i]+1] -=1
　an = np.cumsum(bn)

print(*an)