C++で解く「ボックスから入手できるキャンディーの最大数」問題 ― BFSによる効率的なアプローチ

問題概要

n個の箱があり、各箱は [status, candies, keys, containedBoxes] の形式で以下の情報を持つものとします。

• status[i]：box[i] が開いていれば 1、閉じていれば 0
• candies[i]：box[i] に入っているキャンディーの数
• keys[i]：box[i] 内の鍵で開けられる箱のインデックスの配列
• containedBoxes[i]：box[i] の中に入っている箱のインデックスの配列

まず initialBoxes 配列で指定された箱からスタートします。開いている箱からはすべてのキャンディーを取り出せ、その中の鍵を使って新しい箱を開けることや、見つけた箱を入手することができます。

これらのルールに従って取得できるキャンディーの総数の最大値を求めるのがこの問題です。

入力例

status = [1,0,1,0]
candies = [8,6,5,101]
keys = [[], [], [1], []]
containedBoxes = [[1,2],[3],[],[]]
initialBoxes = [0]

出力例

19

考え方

最初に箱0が与えられます。箱0には8個のキャンディーと、箱1・箱2が入っています。箱1は閉まっていて鍵も持っていないため、まず箱2を開けます。箱2には5個のキャンディーと「箱1の鍵」が入っています。この鍵で箱1を開けると、6個のキャンディーと箱3が見つかりますが、箱3の鍵は存在しないため箱3は閉じたままです。

よって収集できるキャンディーの合計は 8 + 5 + 6 = 19個 となります。

アルゴリズム（幅優先探索）

この問題はキュー（queue）を用いた幅優先探索（BFS）で効率的に解けます。ポイントは「鍵だけ持っていて箱がまだ手元にない状態」「箱はあるが鍵がない状態」を管理することです。以下の手順で進めます。

• ans := 0 と初期化する
• キュー q を用意する
• 集合 visited（手元にある箱）、opened（開封済みの箱）、hasKey（鍵を所持している箱）を定義する
• 初期箱 ib の各要素 x について：
  x を visited に追加し、st[x] == 1（最初から開いている）なら、ans += cnt[x] として x を opened と q に登録する
• q が空になるまで以下を繰り返す：
  先頭要素 curr を取り出し、次の処理を行う。
  • 鍵の処理：k[curr] 内の各 x を hasKey に追加し、x が未開封かつ visited に存在するなら、ans += cnt[x] として x を開けて q に追加する
  • 箱の処理：cb[curr] 内の各 x を visited に追加し、x が未開封かつ（hasKey に存在する または st[x] == 1）なら、x を開けて ans += cnt[x] とし q に追加する
• 最後に ans を返す

C++実装例

それでは、理解を深めるために実際の実装を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
    public:
    int maxCandies(vector<int>& st, vector<int>& cnt,
    vector<vector<int>>& k, vector<vector<int>>& cb, vector<int>& ib) {
        int ans = 0;
        queue<int> q;
        set<int> visited;
        set<int> opened;
        set<int> hasKey;
        for (int i = 0; i < ib.size(); i++) {
            int x = ib[i];
            visited.insert(x);
            if (st[x] == 1) {
                ans += cnt[x];
                opened.insert(x);
                q.push(x);
            }
        }
        while (!q.empty()) {
            int curr = q.front();
            q.pop();
            for (int i = 0; i < k[curr].size(); i++) {
                int x = k[curr][i];
                hasKey.insert(x);
                if (!opened.count(x) && visited.count(x)) {
                    ans += cnt[x];
                    q.push(x);
                    opened.insert(x);
                }
            }
            for (int i = 0; i < cb[curr].size(); i++) {
                int x = cb[curr][i];
                visited.insert(x);
                if (!opened.count(x) && (hasKey.count(x) || st[x] == 1)) {
                    opened.insert(x);
                    ans += cnt[x];
                    q.push(x);
                }
            }
        }
        return ans;
    }
};
main(){
    Solution ob;
    vector<int> v = {1,0,1,0}, v1 = {8,6,5,101}, v2 = {0};
    vector<vector<int>> v3 = {{},{},{1},{}}, v4 = {{1,2},{3},{0},{}};
    cout << (ob.maxCandies(v, v1, v3, v4, v2));
}

入力

{1,0,1,0}, {8,6,5,101}, {{},{},{1},{}}, {{1,2},{3},{0},{}}, {0}

出力

19

まとめ

この問題の本質は、「鍵を先に手に入れて後から箱を得るケース」と「箱を先に手に入れて後から鍵を得るケース」の両方を正しく扱う点にあります。visited・opened・hasKey の3つの集合で状態を管理し、BFSで開封可能な箱を順次処理することで、O(N)程度の計算量で最大キャンディー数を求められます。