概要

この問題は、横一列に並んだ \(N\) 本の花の高さを、隣り合う \(K\) 本の範囲内で「最も高い花と最も低い花の高さの差が \(D\) 以下」という条件を満たしつつ、できるだけ高く残す（高さの総和を最大化する）問題です。

条件をグラフの最短経路問題（ダイクストラ法）に似た形に言い換えることで、 \(O(N \log N)\) の時間計算量で効率よく解くことができます。

---

考察

1. 条件の言い換え

「連続する \(K\) 本の花の中で、最大値と最小値の差が \(D\) 以下」という条件は、以下のように言い換えることができます。

「距離が \(K-1\) 以下の任意の2つの花 \(i, j\) について、高さの差が \(D\) 以下（すなわち \(|H'_i - H'_j| \le D\)）である」

なぜなら、距離が \(K-1\) 以下の2つの花は、必ずどこかの「連続する \(K\) 本の区間」に同時に含まれるため、その差は \(D\) 以下でなければなりません。逆に、距離が \(K-1\) 以下のすべてのペアの差が \(D\) 以下であれば、どの長さ \(K\) の区間を見ても最大値と最小値の差は \(D\) 以下になります。

したがって、各花 \(i\) の最終的な高さ \(H'_i\) は、元の高さ \(H_i\) 以下であることに加えて、自分から距離 \(K-1\) 以内にある他のすべての花 \(j\) の高さ \(+ D\) 以下でなければなりません。

数式で表すと、以下のようになります。 $\(H'_i = \min \left( H_i, \min_{|i-j| \le K-1} (H'_j + D) \right)\)$

2. なぜダイクストラ法のように解けるのか？

上記の式は、グラフの最短経路問題（特にダイクストラ法）における「頂点の最短距離の更新式」と非常に酷似しています。

高さが低い花ほど、周囲の花の高さを「自身の高さ \(+ D\)」以下に強く制限します。 そのため、最終的な高さが低くなる花から順番に、周囲（距離 \(K-1\) 以内）の花へ「高さの制限」を伝播させていくというアプローチが有効です。

具体的には、現在の仮の高さが最も低い花 \(u\) を取り出し、その高さを確定させます。その後、 \(u\) から距離 \(K-1\) 以内にある未確定の花 \(v\) の高さを \(\min(H_v, H'_u + D)\) に更新します。

3. 高速化の工夫（std::set の利用）

愚直に「距離 \(K-1\) 以内のすべての花を探索して更新する」と、 \(K\) が大きい場合に \(O(N \cdot K)\) の時間がかかり、実行時間制限に間に合いません（TLE）。

しかし、ある花 \(v\) が一度でも「他の花からの制限」を受けて高さが更新された（＝確定した）場合、それ以降に別の花からさらに高い制限を受けることはありません。 したがって、まだ高さが確定していない（更新されていない）花のインデックスだけを効率よく見つけることができれば、無駄な探索を省けます。

これは、未確定のインデックスを二分探索木である std::set で管理することで解決できます。 確定した花 \(u\) から距離 \(K-1\) 以内の未確定の花を探す際、set の lower_bound を使うことで、範囲 \([u - (K-1), u + (K-1)]\) に含まれる未確定の花を高速に列挙し、更新と同時に set から削除することができます。

---

アルゴリズム

1. コーナーケースの処理 \(K=1\) のときは、異なる花同士の間に制約はありません。よって、どの花も切る必要はなく、初期状態の高さの総和 \(\sum H_i\) がそのまま答えになります。

2. 初期化

   • \(L = K - 1\) とします。
   • 各花の最終的な高さを格納する配列 \(F\) を用意し、初期値として \(H\) をコピーします。
   • まだ確定していないインデックスを管理する active_indices（std::set）に、すべてのインデックス \(0, 1, \ldots, N-1\) を追加します。
   • 優先度付きキュー pq に、すべての花について (初期の高さ, インデックス) を追加します。

3. ダイクストラ法風の伝播処理 pq が空になるまで以下を繰り返します。

   • pq から最も高さが低い要素 (d, u) を取り出す。
   • すでに u が確定している（訪問済み）ならスキップ。
   • u を訪問済みにし、active_indices から削除する。
   • active_indices の中から、範囲 \([u - L, u + L]\) に含まれるインデックス \(v\) を lower_bound で探索する。
   • 条件を満たす各 \(v\) について：
     • \(F_v\) を \(\min(H_v, d + D)\) に更新する。
     • 更新された (F[v], v) を pq に追加する。
     • \(v\) を確定したとみなし、active_indices から削除する。

4. 答えの出力 最終的な配列 \(F\) の要素の総和を計算して出力します。

---

計算量

• 時間計算量: \(O(N \log N)\)

  • 各要素は active_indices（std::set）に高々1回しか追加されず、1回しか削除されません。
  • lower_bound による探索と削除の操作は、全体を通して高々 \(N\) 回しか実行されません（アモルタズ解析）。
  • 優先度付きキューへの追加・取り出し操作も、各要素について高々数回（初期化時と更新時）しか行われないため、全体で \(O(N \log N)\) となり、十分に高速です。

• 空間計算量: \(O(N)\)

  • 配列 \(H, F\)、訪問済みフラグ、std::set、および優先度付きキューのサイズはいずれも最大で \(N\) に比例するため、メモリ制限にも余裕で間に合います。

---

実装のポイント

• std::set からの効率的な削除 C++において、std::set から範囲内の要素をループで削除しながらイテレータを進める際、it = active_indices.erase(it) と書くことで、要素を削除しつつ、次の要素を指す有効なイテレータを \(O(1)\)（ならし時間）で取得できます。

• オーバーフローの防止 花の高さやその総和、および許容差 \(D\) は非常に大きな値（最大で \(10^9\) や \(N \times 10^9\)）になるため、C++では long long 型（64bit整数型）を使用する必要があります。

  ソースコード

#include <iostream>
#include <vector>
#include <queue>
#include <set>
#include <algorithm>

using namespace std;

int main() {
    // Optimize standard I/O operations for competitive programming
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int N, K;
    long long D;
    if (!(cin >> N >> K >> D)) return 0;

    vector<long long> H(N);
    long long sum_H = 0;
    for (int i = 0; i < N; ++i) {
        cin >> H[i];
        sum_H += H[i];
    }

    // Corner case: if K = 1, there are no constraints between different flowers.
    if (K == 1) {
        cout << sum_H << "\n";
        return 0;
    }

    long long L = K - 1;

    vector<long long> F(N);
    vector<bool> visited(N, false);
    set<int> active_indices;
    using P = pair<long long, int>;
    priority_queue<P, vector<P>, greater<P>> pq;

    for (int i = 0; i < N; ++i) {
        active_indices.insert(i);
        pq.push({H[i], i});
        F[i] = H[i];
    }

    while (!pq.empty()) {
        auto [d, u] = pq.top();
        pq.pop();

        if (visited[u]) continue;
        visited[u] = true;

        if (active_indices.count(u)) {
            active_indices.erase(u);
        }

        long long l = max(0LL, u - L);
        long long r = min((long long)N - 1, u + L);

        // Find all active indices in the range [u - L, u + L]
        auto it = active_indices.lower_bound(l);
        while (it != active_indices.end() && *it <= r) {
            int v = *it;
            if (H[v] > d + D) {
                F[v] = d + D;
                pq.push({F[v], v});
            } else {
                F[v] = H[v];
                pq.push({F[v], v});
            }
            // Erase from set and get the next iterator in O(1) amortized time
            it = active_indices.erase(it);
        }
    }

    long long ans = 0;
    for (int i = 0; i < N; ++i) {
        ans += F[i];
    }

    cout << ans << "\n";

    return 0;
}

---

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