公式

B - フルーツの詰め合わせ / Fruit Assortment 解説 by MMNMM


与えられた \(N+M\) 個の整数の大きいほうから \(K\) 個を選び、それらの合計を求めればよいです。

  • ソートを使うことで \(O((N+M)\log (N+M))\) 時間
  • ヒープを使うことで \(O((N+M)\log K)\) 時間
  • Quick Select を使うことで \(O(N+M)\) 時間

などで実現できます。

実装例は以下のようになります。

#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 + M);
    for (int& a : A) {
        cin >> a;
    }

    // O(N) 時間で大きいほうから K 個を先頭に移動させる
    ranges::nth_element(A, begin(A) + K, greater{});
    // 先頭 K 個の合計を求める
    cout << ranges::fold_left(A | views::take(K), 0UL, plus{}) << endl;
    return 0;
}
N, M, K = map(int, input().split())

A = [int(input()) for i in range(N + M)]

# 降順にソートして、先頭 K 個の合計を求める
print(sum(sorted(A, reverse=True)[:K]))

投稿日時:
最終更新: