C++で解く「森のウサギ」問題 ― 最小のウサギ数を求めるアルゴリズム

問題の概要

森の中にいるすべてのウサギには、それぞれ固有の毛色があります。ここで、一部のウサギ(全員の場合もありえます)が「自分と同じ色のウサギは他に何匹いるか」という質問に答え、その回答が配列として与えられます。求めたいのは、これらの回答と矛盾しない範囲で、森に存在しうるウサギの最小数です。

入力例での考え方

たとえば入力が [1, 1, 2] のとき、答えは 5 になります。

• 「1」と答えた2匹は、同じ白色のグループに属していると考えられます(自分のほかに同じ色のウサギが1匹いる、という意味です)。
• 一方、「2」と答えたウサギは白ではあり得ません。もし白なら「同じ色は1匹」と答えるはずだからです。そこでこのウサギを黒色と仮定すると、配列に回答しなかった黒いウサギがあと2匹森にいる必要があります。

以上より、回答した3匹と回答しなかった2匹を合わせて、森のウサギの最小数は 5匹 となります。

解法のアプローチ

鍵となるのは次の観察です。同じ数値 x と答えたウサギ同士は、最大 x+1 匹まで同じグループ(同じ色)にまとめられるということです。グループが満員になった後に同じ数値の回答が現れた場合は、それは別の新しいグループとして扱う必要があります。

この性質を利用し、以下の手順で貪欲にグループ化を行います。

1. マップ m を用意し、n を配列 ans のサイズ、ret を 0 で初期化します。
2. i を 0 から n−1 まで順に処理します。
   • x := ans[i] とします。
   • x == 0 の場合:その色のウサギは自分1匹だけなので、ret を 1 増やして次の反復へ進みます。
   • m に x が存在しない場合:新しいグループを作るため、ret に (x + 1) を加算し、m[x] := 0 とします。
   • m に x がすでに存在する場合:m[x] を 1 増やします。そして m[x] == x になったら、そのグループは満員なので m から x を削除します。
3. 最後に ret を返します。

C++による実装

理解を深めるために、以下の実装例を見てみましょう。

#include <bits/stdc++.h>
using namespace std;
class Solution {
public:
   int numRabbits(vector<int>& ans) {
      map <int, int> m;
      int n = ans.size();
      int ret = 0;
      for(int i = 0; i < n; i++){
         int x = ans[i];
         if(x == 0){
            ret++;
            continue;
         }
         if(!m.count(x)){
            ret += (x + 1);
            m[x] = 0;
         }else{
            m[x]++;
            if(m[x] == x){
               m.erase(x);
            }
         }
      }
      return ret;
   }
};
main(){
   vector<int> v = {1,1,2};
   Solution ob;
   cout << (ob.numRabbits(v));
}

入力

[1,1,2]

出力

5

動作のポイントと計算量

このアルゴリズムでは、同じ回答を持つウサギをできる限り1つのグループに詰め込むことで、回答しなかった(見えていない)ウサギの数を最小化しています。回答 x ごとに必要なグループサイズは x+1 匹であり、グループが満員になるまでは追加コストなしにウサギを受け入れられるのがポイントです。

• 時間計算量: O(n log n)(std::map の操作コストを含む。unordered_map を使えば平均 O(n) にできます)
• 空間計算量: O(n)