C++で解く「バスケットに果物を入れる」問題 ― スライディングウィンドウによる実装

一列に並んだ木があり、i 番目の木には tree[i] という種類の果実が実っているとします。私たちは好きな木からスタートし、以下の手順を繰り返し実行します。

• 今いる木から果実を 1 つバスケットに加えます。果実を追加できない場合は終了します。
• 現在の木の一つ右隣の木へ移動します。右側に木が存在しない場合は終了します。

バスケットは 2 つあり、それぞれどれだけの量でも運ぶことができますが、1 つのバスケットに入れられる果実は 1 種類だけという制約があります。この手順で収集できる果実の合計数を求めましょう。

例えば、木が [0, 1, 2, 2] のように並んでいる場合、答えは 3 になります。2 本目以降の木からスタートすれば [1, 2, 2] を収集できますが、最初の木から始めると [0, 1] しか集められないためです。

アプローチ

この問題は「スライディングウィンドウ（尺取り法）」を用いることで効率的に解けます。ウィンドウ内に含まれる果実の種類が常に 2 種類以下になるよう区間を伸縮させるのがポイントです。手順は以下の通りです。

• n := 木の本数、j := 0、ans := 0 として初期化します。
• マップ m を作成します。
• i を 0 から n − 1 までループさせます。
  • m[tree[i]] を 1 増やします。
  • m のサイズが 2 より大きく、かつ j ≤ i の間、次を繰り返します。
    • m[tree[j]] を 1 減らします。
    • m[tree[j]] == 0 になったら、m から tree[j] を削除します。
    • j を 1 増やします。
  • ans := max(i − j + 1, ans) と更新します。
• ans を返します。

このアルゴリズムでは、左端 j と右端 i がどちらも後戻りせずに前へ進むため、配列全体を一度の走査で処理できます。時間計算量は O(n) となり、またマップに保持される果実の種類は高々 3 種類なので、空間計算量は O(1) と非常に効率的です。

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

C++ 実装例

#include <bits/stdc++.h>
using namespace std;
class Solution {
   public:
   int totalFruit(vector<int>& tree) {
      int n = tree.size();
      int j = 0;
      map <int, int> m;
      int ans = 0;
      for(int i = 0; i < n; i++){
         m[tree[i]] += 1;
         while(m.size() > 2 && j <= i){
            m[tree[j]]--;
            if(m[tree[j]] == 0)m.erase(tree[j]);
            j++;
         }
         ans = max(i - j + 1, ans);
      }
      return ans;
   }
};
main(){
   vector<int> v = {3,3,3,1,2,1,1,2,3,3,4};
   Solution ob;
   cout <<(ob.totalFruit(v));
}

入力

[3,3,3,1,2,1,1,2,3,3,4]

出力

5

この入力の場合、「1 種類または 2 種類の果実のみを含む連続区間」の中で最も長くなるのは [1, 2, 1, 1, 2] の 5 個であるため、出力は 5 になります。