問題概要
この問題は単純に線形探索で解こうとすると2重ループになり、実行時間制限の２秒を超えてしまう。
そのためどうにかして2重ループを避ける必要がある。
愚直解
まずは愚直解。というか私は最初この方法で解いてしまった。
N, M = map(int, input().split())
A = list(map(int, input().split()))
# 後続の処理でAの要素（花火が上がる日）を探索したいためlistをsetに変換
hanabi_day_set = set()
for i in range(M):
hanabi_day_set.add(A[i])
# 花火が上がる日をTrue, そうでない日をFalseをセット
day_flg_list = [False] * (N + 1)
for i in range(1, N + 1):
if i in hanabi_day_set:
day_flg_list[i] = True
# １日目からN日目まで線形探索して、i日目から何日後に花火が上がるか確認する
for i in range(1, N + 1):
count = 0
for j in range(i, N + 1):
if day_flg_list[j]:
print(count)
break
else:
count += 1
この方法だとforの二重ループを使っているため、計算量はO(N^2)かかる。このように単純に実装してしまったらTLEになる。
改善案
計算量O(N)でどうやって実行できるか考えてみる。
愚直解の外側のループについては、N個の要素に対して花火は上がるのか、上がらないのかを確認する必要があるため、残す必要がある。
よって内側のループをもう少し最適化できないか考えてみる。
見直してみると、愚直解の内側のループは各要素から次の要素までどれくらい離れているか要素を一つずつ辿っていっている。実はここを一つ前の結果を使ってその日が次花火が上がる日からどれくらい離れているかを求めることができる。
今回は後ろから要素を見ていき、もし花火が上がる日であれば0を追加、もしそうでなければ前の要素+1日を足した要素を追加する。そして最後に逆順に並び替える。
そうすることで各要素を辿らず、O(1)で計算することができる。
以下、解答例。
N, M = map(int, input().split())
A = list(map(int, input().split()))
hanabi_day_set = set()
for i in range(M):
hanabi_day_set.add(A[i])
day_flg_list = [False] * (N + 1)
for i in range(1, N + 1):
if i in hanabi_day_set:
day_flg_list[i] = True
# 後ろの要素から見ていって、花火が上がる日であれば0、そうでなければ最後に確認した要素から1足した値を設定する
ans_list = [0]
for i in range(N - 1, 0, -1):
if day_flg_list[i]:
ans_list.append(0)
else:
ans_list.append(ans_list[N - i - 1] + 1)
# 最後に逆順に並び替えて、出力する
rev_list = list(reversed(ans_list))
for i in range(N):
print(rev_list[i])
参考
別解
この方法はある日にち（1日目、2日目…）を基準にして、その日にちがどれくらい次花火が上がる日から離れているのか計算している。
この方法も内側のループはO(1)で実行可能である。
N, M = map(int, input().split())
A = list(map(int, input().split()))
day = 1
for i in range(M):
while A[i] - day >= 0:
print(A[i] - day)
day += 1
まとめ
今回やってしまった愚直解は、要素を一つずつ辿ってどれくらい離れているかを確認していたのがよくありませんでした。まあ今回のような問題で配列の要素を一つずつ見ていくようではそりゃTLEになりますよね。。
今回の問題から前の計算結果を使って次の値を計算することで、計算量をO(1)にすることを頭の片隅に入れておきたいと思いました。