東京海上日動 プログラミングコンテスト2020 C - Lamps
備忘録
問題
回答
import sys import os import math import bisect import itertools import collections import heapq import queue import array import time # 時々使う # from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall # from decimal import Decimal # from collections import defaultdict, deque # 再帰の制限設定 sys.setrecursionlimit(10000000) def ii(): return int(sys.stdin.buffer.readline().rstrip()) def il(): return list(map(int, sys.stdin.buffer.readline().split())) def fl(): return list(map(float, sys.stdin.buffer.readline().split())) def iln(n): return [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)] def iss(): return sys.stdin.buffer.readline().decode().rstrip() def sl(): return list(map(str, sys.stdin.buffer.readline().decode().split())) def isn(n): return [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)] def lcm(x, y): return (x * y) // math.gcd(x, y) MOD = 10 ** 9 + 7 INF = float('inf') def main(): if os.getenv("LOCAL"): sys.stdin = open("input.txt", "r") N, K = il() A = il() for k in range(K): acc = [0] * N for n in range(N): l = max(0, n-A[n]) r = n+A[n]+1 acc[l] += 1 if r < N: acc[r] -= 1 A = list(itertools.accumulate(acc)) if sum(A) == N ** 2: break print(*A) if __name__ == '__main__': main()
考え方
累積和を用いて、電球が照らす範囲を更新します。
なお、上記の回答例はPyPy3 (7.3.0)
で提出しています。
- 回答概要
- 累積和を用いて電球が照らす範囲を更新する
N
個の電球が全ての座標を照らした時点で処理を終了する
光の強さAi
の電球は、電球の位置i
と電球の位置i
から左右にAi
だけ進んだ位置までを照らすことが出来ます。
数直線を配列で用意し、電球の照らす範囲を愚直に全て更新した場合、到底2秒では間に合わずLTEとなります。
そのため、累積和を用いて電球が照らす範囲の更新を行います。
具体的には、位置i
に存在する電球の光の強さがAi
のとき、
電球が照らすことができる最も左の位置(max(0, i-A[i])
)に+1
し、
電球が照らすことが出来ない右の位置(n+A[n]+1
)に-1
することで、
累積和を求めたとき、位置i
にいくつの電球が照らしているか求めることが出来ます。
ただし、注意点として、照らす範囲の更新をK
回行うとき、
K
は最大でも2 * 10 ** 5
回行われます。
そのため、素直にK
回だけ更新処理を行うとLTEとなります。
LTEを防ぐために、全ての位置i
に対して光を照らす電球の数が変わらなくなって時点で、
更新処理を終了させる必要があります。
1回の更新処理で全ての位置に対して光を当てる電球は最低でも1つずつ増えます。
つまり、ある回数以上行うと、必ず全ての位置において、光を当てる電球の数はN
個になります。
それ以上処理を行っても電球の数は更新されないため、その時点で処理を終了させる必要があります。