【AtCoder】ABC 358 D - Souvenirs | 茶コーダーが解くAtCoder

atcoder.jp

実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 393

問題概要

お土産屋に の番号が付いた 個の箱が売られている。箱 の価格は 円であり、箱 の中には 個のお菓子が入っている。高橋君は 個の箱のうち 個を選んで買って帰り、 の名前が付いた 人に 1 つずつ箱を渡そうと考えている。ただし、以下の条件を満たすことができるように箱を買いたい。

• 各 について、人 には 個以上のお菓子が入った箱を渡す。

1人に2個以上の箱を渡したり、同じ箱を複数人に渡したりすることはできない。これらの条件を満たしつつ適切に 個の箱を買うことができるか判定し、可能な場合は支払う金額の合計の最小値を出力せよ。そうでない場合は-1を出力せよ。

制約

• 入力はすべて整数。
• 
• 

multisetを用いた解法

考察

やりたいこと自体は単純で、「各 について、 個の箱の中でお菓子が 個以上入っているものを抽出し、その中でお菓子の個数が最小のものを購入する」という貪欲法が最適である。本問では箱に入っているお菓子の個数と価格が一致していることに注目しよう。この証明については公式解説を参照のこと。

これを実現するためには、

1. 各箱のお菓子の個数を格納した配列の中で を探す。
   1. 見つかればそれを購入する。
   2. 見つからなければその時点で-1を出力して処理を打ち切る。
2. 購入した箱は配列から削除する。

とするのが良いだろう。

ここで、手順1における探索と手順4における要素の削除については、長さ のvectorに対して愚直にループで行ったりerase()で行ったりすると、どちらも計算量が になってしまう点に注意しよう。本問の制約から、これは TLE になってしまう。

C++ では、各箱のお菓子の個数を管理するデータ構造としてmultisetを採用し、二分探索 (lower_bound) を用いることでこの問題を回避できる。

cpprefjp.github.io

また、multisetに対して要素を一つだけ削除する場合は、eraseの引数に要素のイテレータを指定することに注意。そのまま値を指定すると、対応する要素全てが削除されてしまう。

コードではBをソートしているが、必ずしも必要なわけではない。

コード

#include <bits/stdc++.h>
using namespace std;

#define all(x) (x).begin(), (x).end()
#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using ll = long long;

// ======================================== //

int main()
{
    int N, M;
    cin >> N >> M;
    multiset<ll> st;
    vector<ll> B(M);
    rep(i, 0, N)
    {
        ll A;
        cin >> A;
        st.insert(A);
    }
    rep(i, 0, M) cin >> B[i];

    sort(all(B));

    ll ans = 0;
    rep(i, 0, M)
    {
        auto itr = st.lower_bound(B[i]);
        if (itr == st.end())
        {
            cout << -1 << endl;
            return 0;
        }
        ans += *itr;
        st.erase(itr);
    }

    cout << ans << endl;
}

atcoder.jp

実装時間: 25分

priority_queueを用いた解法

前述したmultisetは使用頻度が比較的少ないので、思いつかなかった人もいるかもしれない (私もです) 。その場合でも、priority_queueを用いることで以下のような実装ができる。ただし、やや冗長である。

考察

ここでは、要求するお菓子の数が多いものから考えていく (小さい順に考えても、ソートの向きを逆にすることで対応可能なはず) 。

まず、箱の中のお菓子の個数と価格を管理する構造体Boxを作成し、 の情報からBox型の配列boxesを作成。これを独自の比較関数compareを用いてお菓子の数が多い順にソートする。

続けて、 も降順にソートする。

以降は、価格が安い順にソートされる優先度付きキューであるqueを作成し、各 に対して順に見ていく。

1. 「 かつ のお菓子の数が 以上である」という条件を満たす を から順に探していき、queに追加していく。
2. すべて追加し終わった後、queの先頭 (条件を満たす中で最も価格が安いもの) をpopし、合計金額に加える。
3. もしqueが空ならば、その時点で-1を出力して処理を打ち切る。

次の に移る過程で はリセットされないので、購入済みの箱は手順1における探索からは除外される。

コード

#include <bits/stdc++.h>
using namespace std;

#define all(x) (x).begin(), (x).end()
#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

using ll = long long;

// ======================================== //

struct Box
{
    int price;
    int snacks;
};

bool compare(const Box &a, const Box &b)
{
    return a.snacks > b.snacks;
}

int main()
{
    int N, M;
    cin >> N >> M;
    vector<int> A(N);
    vector<int> B(M);
    rep(i, 0, N) cin >> A[i];
    rep(i, 0, M) cin >> B[i];

    vector<Box> boxes(N);
    rep(i, 0, N) boxes[i] = {A[i], A[i]};

    sort(all(boxes), compare);
    sort(all(B), greater<int>());

    int j = 0;
    ll sum = 0;
    priority_queue<int, vector<int>, greater<int>> que;
    rep(i, 0, M)
    {
        while (j < N && boxes[j].snacks >= B[i])
        {
            que.push(boxes[j].price);
            j++;
        }

        if (!que.empty())
        {
            sum += que.top();
            que.pop();
        }
        else
        {
            cout << -1 << endl;
            return 0;
        }
    }

    cout << sum << endl;

    return 0;
}

atcoder.jp

実装時間: 35分