概要

高橋君が決められたパターンでリンゴを収穫する際に、収穫するリンゴの総数を求める問題です。

考察

問題では、高橋君は連続する \(K\) 本の木からリンゴを収穫した後、次の1本の木をスキップするというパターンを繰り返します。このパターンは木1から始まり、木Nに到達するまで続けられます。素朴なアプローチとして、whileループで現在位置を管理し、各ステップで \(K\) 本分のリンゴを加算し、スキップ分を飛ばすことで、効率的に総和を計算できます。入力サイズが最大 \(10^6\) と大きいため、単純なシミュレーションでも \(O(N)\) で処理可能です。

アルゴリズム

1. 現在の位置 idx を0（0-indexed）で初期化する。
2. idx が \(N\) 未満の間、以下を繰り返す：
   • 現在位置から \(K\) 本分の木（ただし残りが少ない場合は残り全て）のリンゴの個数を総和に加算する。
   • 現在位置を \(K\) 本分進めた後、さらに1本スキップする（idx = end + 1）。
3. 総和を出力する。

計算量

• 時間計算量: \(O(N)\)
  • 各リンゴは高々1回しか加算されないため。
• 空間計算量: \(O(N)\)
  • リンゴの個数を格納する配列のサイズが \(N\) であるため。

実装のポイント

• 0-indexedで配列を扱うため、初期位置を0としている。

• min(idx + K, N) で、残りの木が \(K\) 本未満の場合の処理を簡潔に表現している。

• ループ内で直接加算するため、余計なデータ構造を使わずシンプルに実装している。

  ソースコード

def main():
    import sys
    data = sys.stdin.read().split()
    N = int(data[0])
    K = int(data[1])
    D = list(map(int, data[2:2+N]))
    
    total = 0
    idx = 0
    while idx < N:
        end = min(idx + K, N)
        for i in range(idx, end):
            total += D[i]
        idx = end + 1
    print(total)

if __name__ == "__main__":
    main()

この解説は deepseekv3 によって生成されました。