概要

\(N\) 個の街灯があり、特定の街灯の電球を交換すると、その街灯自身と両隣の街灯の明るさが \(1\) ずつ増加します。\(M\) 回の交換作業を行った後の、各街灯の最終的な明るさを求める問題です。

考察

この問題のポイントは、「ある街灯 \(i\) の明るさが、どの交換作業によって増加するか」を整理することです。

街灯 \(i\) の明るさが増加するのは、以下のいずれかの街灯の電球が交換されたときです。 1. 街灯 \(i-1\) の電球が交換されたとき（右隣である \(i\) に光が届く） 2. 街灯 \(i\) の電球が交換されたとき（自分自身） 3. 街灯 \(i+1\) の電球が交換されたとき（左隣である \(i\) に光が届く）

したがって、すべての作業が終わった後の街灯 \(i\) の明るさは、以下の合計になります。 $\((\text{初期の明るさ } A_i) + (\text{街灯 } i-1 \text{ が交換された回数}) + (\text{街灯 } i \text{ が交換された回数}) + (\text{街灯 } i+1 \text{ が交換された回数})\)$

素朴なシミュレーションとして、「交換作業 \(B_j\) が発生するたびに、隣接する 3 つの街灯の明るさを直接更新する」という方法も考えられます。この場合、1回の作業につき最大 3 箇所の更新を行うため、全体の計算量は \(O(M + N)\) となり、今回の制約（\(N, M \leq 2 \times 10^5\)）では十分に間に合います。

提供されたコードでは、より整理されたアプローチとして、まず各街灯が何回交換されたかをカウントし、最後に各街灯の最終的な明るさを一気に計算しています。

アルゴリズム

以下の手順で解いています。

1. 交換回数のカウント: 長さ \(N+2\) の配列 count を用意します（1-indexed で管理し、端の処理を楽にするために少し大きめに確保します）。 与えられた \(B_1, B_2, \ldots, B_M\) を走査し、count[B_j] を \(1\) ずつ増やします。

2. 最終的な明るさの計算: 各街灯 \(i = 1, 2, \ldots, N\) について、以下の計算を行います。 $\(\text{ans}_i = A_i + \text{count}[i-1] + \text{count}[i] + \text{count}[i+1]\)\( ここで、街灯 \)1\( のとき `count[0]`、街灯 \)N\( のとき `count[N+1]` を参照しますが、これらは \)0$ のままなので、条件分岐なしで正しく計算できます。

3. 出力: 計算した結果をスペース区切りで出力します。

計算量

• 時間計算量: \(O(N + M)\)
  • 入力の読み込みに \(O(N + M)\)、交換回数のカウントに \(O(M)\)、最終的な明るさの計算に \(O(N)\) かかります。
• 空間計算量: \(O(N + M)\)
  • 入力データ、明るさの配列 \(A\)、交換回数の配列 count を保持するために必要です。

実装のポイント

• 高速な入出力: Python では \(N, M\) が \(2 \times 10^5\) 程度になると、標準の input() を繰り返すと実行時間制限（TLE）になる恐れがあります。sys.stdin.read().split() を使って一括で読み込むことで高速化しています。

• 番兵（境界条件の簡略化）: count 配列のサイズを \(N+2\) にすることで、街灯 \(1\) の左隣（インデックス 0）や街灯 \(N\) の右隣（インデックス \(N+1\)）へのアクセスをエラーにせず、かつ値 \(0\) として扱えるようにしています。これにより、if 文による端の処理を省いています。

  ソースコード

import sys

def solve():
    # 読み込みの高速化のため、すべての入力を一括で取得して分割します
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    # N: 街灯の数, M: 電球交換の回数
    N = int(input_data[0])
    M = int(input_data[1])
    
    # A: 各街灯の初期の明るさ（0-indexedのリストとして格納）
    # input_data[2] から input_data[2+N-1] までが A_1 から A_N に対応
    A = list(map(int, input_data[2 : 2 + N]))
    
    # count[i]: 街灯 i の電球が交換された回数
    # 1-indexedで管理するため、サイズを N+2 とし、境界条件（0 と N+1）もカバーします
    count = [0] * (N + 2)
    # input_data[2+N] から input_data[2+N+M-1] までが B_1 から B_M に対応
    for i in range(2 + N, 2 + N + M):
        b = int(input_data[i])
        count[b] += 1
        
    # すべての作業が終わった後の各街灯の明るさを計算します
    # 街灯 i の明るさは、初期の明るさ A_i に加え、
    # 街灯 i-1, i, i+1 の電球が交換された回数の合計分だけ増加します
    ans = [None] * N
    for i in range(1, N + 1):
        # A[i-1] は街灯 i の初期の明るさ
        # count[i-1], count[i], count[i+1] はそれぞれ隣接および自身の電球交換回数
        # count[0] と count[N+1] は常に 0 であるため、端の街灯も正しく計算されます
        final_brightness = A[i-1] + count[i-1] + count[i] + count[i+1]
        ans[i-1] = str(final_brightness)
        
    # 結果をスペース区切りで1行に出力します
    sys.stdout.write(" ".join(ans) + "\n")

if __name__ == '__main__':
    solve()

この解説は gemini-3-flash-preview によって生成されました。