Official

B - お菓子の詰め合わせ / Assortment of Sweets Editorial by MMNMM


この制約のもと、ありえるセットは最大で \(4498500\) 個です。 それぞれのセットの満足度を求め、上位 \(M\) 個の満足度の合計を求めればよいです。

セットの満足度を求める際、それぞれのセットに対して毎回おいしさの合計を求めると最悪時間計算量は \(\Theta(NK ^ 2)\) となってしまいます。 累積和を用いるか、区間を伸ばしながら合計を求めていくことで全体の計算量を \(O(NK)\) とすることができます。

上位 \(M\) 個の合計を求めるアルゴリズムは以下のような選択肢があります。

  • すべての満足度を求めてソートし、上位 \(M\) 個の合計を計算した場合、時間計算量は \(O(NK\log(NK))\)
  • ヒープなどを使い、上位 \(M\) 個を保持しながらすべての満足度を求めて合計した場合、時間計算量は \(O(NK\log M)\)
  • 適切な選択アルゴリズム(Quick Select など)を用いた場合、時間計算量は \(O(NK)\)

これらのどれを採用しても十分高速です。

実装例は以下のようになります。 \(1\) つのセットに対する満足度の値が \(32\) bit 符号付き整数型の範囲に収まらない場合があることに注意してください。

#include <iostream>
#include <vector>
#include <algorithm>
#include <ranges>
using namespace std;

int main() {
    int N, M, K;
    cin >> N >> M >> K;

    vector<int> A(N);
    for (int& a : A) {
        cin >> a;
    }

    // ありえるセットすべてに対して満足度を求める
    vector<long> set_satisfaction;
    for (int i = 0; i < N; ++i) { // 始点を決めて
        long s = 0;
        for (int k = 0; k < K && i + k < N; ++k) { // 区間を伸ばしながら合計を計算
            s += A[i + k];
            set_satisfaction.push_back(s);
        }
    }

    // 降順に並べた先頭 M 個を求める
    ranges::nth_element(set_satisfaction, set_satisfaction.begin() + M, greater{});

    // 合計を出力
    cout << ranges::fold_left(set_satisfaction | views::take(M), 0L, plus{}) << endl;
    return 0;
}
N, M, K = map(int, input().split())
A = list(map(int, input().split()))

# ありえるセットすべてに対して満足度を求める
set_satisfaction = []
for i in range(N): # 始点を決めて
    s = 0
    for _, a in zip(range(K), A[i:]): # 区間を伸ばしながら合計を計算
        s += a
        set_satisfaction.append(s)

print(sum(sorted(set_satisfaction, reverse=True)[:M]))

posted:
last update: