概要

\(N\) 本の木から「\(K\) 本収穫して \(1\) 本休む」というパターンでリンゴを収穫したときの、収穫したリンゴの合計を求める問題です。

考察

高橋君の行動を観察すると、「\(K\) 本収穫して \(1\) 本休む」というパターンを繰り返しています。これを言い換えると、高橋君が収穫しない（休む）木は、最初から数えて \(K+1\) 番目、\(2(K+1)\) 番目、\(3(K+1)\) 番目…… となります。

素朴に先頭から \(1\) 本ずつ「収穫する・しない」を判定して足し合わせるループ処理でも解くことができます。しかし、Pythonなどの一部の言語では、さらに簡潔かつ高速に計算する方法があります。 それは、「すべてのリンゴの合計」から「収穫しなかった木のリンゴの合計」を引くというアプローチです。

アルゴリズム

配列のインデックス（\(0\) 始まり）で考えます。 木 \(1\) はインデックス 0、最初に休む木 \(K+1\) はインデックス K となります。 したがって、収穫しない木はインデックス K, 2K+1, 3K+2…… つまり、インデックス K から始まり、K+1 ずつ増加していくことになります。

Pythonのスライス機能 [start:stop:step] を用いると、収穫しない木のリストは D[K::K+1] と一行で簡潔に取得できます。 よって、以下の手順で答えを求めます。 1. すべての木のリンゴの総和 sum(D) を計算する。 2. 収穫しない木のリンゴの総和 sum(D[K::K+1]) を計算する。 3. 1から2を引いた値を答えとする。

計算量

• 時間計算量: \(O(N)\) 配列の要素の合計を求める処理やスライス操作は配列の長さに比例するため、\(O(N)\) で完了します。制約の \(N \le 10^6\) に対して十分高速です。
• 空間計算量: \(O(N)\) 入力を配列として保持するため、\(O(N)\) のメモリを使用します。

実装のポイント

• Pythonのスライスの活用: D[K::K+1] と書くことで、for 文を使って明示的にループを書くよりも、内部的にC言語レベルで最適化されて実行されるため非常に高速に処理されます。

• オーバーフローへの注意: \(D_i\) の最大値が \(10^9\)、\(N\) が \(10^6\) なので、答えは最大 \(10^{15}\) になります。C++などの言語では32bit整数（int）に収まらないため64bit整数（long long）を使う必要がありますが、Pythonでは自動的に多倍長整数として扱われるためオーバーフローの心配はありません。

  ソースコード

import sys

def main():
    input = sys.stdin.read
    data = input().split()
    if not data:
        return
    N = int(data[0])
    K = int(data[1])
    D = [int(x) for x in data[2:]]
    
    ans = sum(D) - sum(D[K::K+1])
    print(ans)

if __name__ == '__main__':
    main()

この解説は gemini-3.1-pro-thinking によって生成されました。