C++で文字列内の全単語を連結した部分文字列の開始インデックスを求める

問題概要

文字列 s と、すべて同じ長さの単語からなるリスト words が与えられたとします。このとき、s の中に存在する部分文字列のうち、words に含まれる各単語をちょうど1回ずつ、間に余計な文字を挟まずに連結したものに一致するものを探し、その開始インデックスをすべて求めます。

たとえば、入力が "barfoothefoobarman"、単語リストが ["foo", "bar"] の場合、出力は [0, 9] となります。これは、インデックス 0 から始まる部分文字列が "barfoo"、インデックス 9 から始まる部分文字列が "foobar" であり、いずれも "bar" と "foo" を1回ずつ連結した形になっているためです。

解法のアプローチ

この問題は、スライディングウィンドウとハッシュマップを組み合わせることで効率的に解けます。手順は以下の通りです。

補助関数 ok() の定義

• 文字列 s、マップ wordCnt、整数 n(1単語の長さ)を受け取ります。
• s の先頭 n 文字を temp にコピーします。
• i を n から s.size() - 1 まで順に処理します。
  • temp の長さが n の倍数になったタイミングで判定を行います。
  • temp が wordCnt に存在しなければ false を返します。
  • 存在する場合、wordCnt[temp] が 1 ならマップから削除し、それ以外ならカウントを 1 減らします。いずれの場合も temp を空文字列に戻します。
  • その後、temp += s[i] で次の文字を追加します。
• ループ終了後、残った temp に対しても同様の判定を行います。
• 最終的に wordCnt が空になっていれば true を返します。

メイン関数 findSubstring() の流れ

• a または b が空の場合は空の配列を返します。
• マップ wordCnt を作成し、b 内の各単語の出現回数を記録します。
• 結果格納用の配列 ans を用意します。
• ウィンドウ幅 window を「単語数 × 1単語あたりの文字数」として計算します。
• 先頭 window 文字を temp に設定し、i を window から a.size() - 1 までスライドさせます。
  • temp の長さが window の倍数で、かつ ok() が true を返せば、開始位置 i - window を ans に追加します。
  • temp の末尾に a[i] を追加し、長さが window を超えたら先頭の1文字を削除します。
• 最後のウィンドウについても同様に判定し、該当すれば a.size() - window を追加します。
• ans を返します。

C++ 実装例

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

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

void print_vector(vector<auto> v){
   cout << "[";
   for(int i = 0; i<v.size(); i++){
      cout << v[i] << ", ";
   }
   cout << "]" << endl;
}

class Solution {
public:
   bool ok(string s, unordered_map <string, int> wordCnt, int n){
      string temp = "";
      for(int i = 0; i < n; i++){
         temp += s[i];
      }
      for(int i = n; i < s.size(); i++){
         if(temp.size() % n == 0){
            if(wordCnt.find(temp) == wordCnt.end())return false;
            else{
               if(wordCnt[temp] == 1){
                  wordCnt.erase(temp);
                  temp = "";
               }
               else{
                  wordCnt[temp]--;
                  temp = "";
               }
            }
         }
         temp += s[i];
      }
      if(wordCnt.find(temp) == wordCnt.end())return false;
      else{
         if(wordCnt[temp] == 1){
            wordCnt.erase(temp);
            temp = "";
         }
         else{
            wordCnt[temp]--;
            temp = "";
         }
      }
      return wordCnt.size() == 0;
   }
   vector<int> findSubstring(string a, vector<string> &b) {
      if(a.size() == 0 || b.size() == 0)return {};
      unordered_map <string, int> wordCnt;
      for(int i = 0; i < b.size(); i++)wordCnt[b[i]]++;
      vector <int> ans;
      int window = b.size() * b[0].size();
      string temp ="";
      for(int i = 0; i < window; i++)temp += a[i];
      for(int i = window; i < a.size(); i++){
         if(temp.size() % window == 0 && ok(temp, wordCnt, b[0].size())){
            ans.push_back(i - window);
         }
         temp += a[i];
         if(temp.size() > window)temp.erase(0, 1);
      }
      if(temp.size() % window == 0 && ok(temp, wordCnt, b[0].size()))ans.push_back(a.size() - window);
      return ans;
   }
};

main(){
   vector<string> v = {"foo", "bar"};
   Solution ob;
   print_vector(ob.findSubstring("barfoothefoobarman", v));
}

実行結果

上記のコードを実行すると、以下の出力が得られます。

[0, 9]

この結果は、文字列 "barfoothefoobarman" の中で、インデックス 0 に "barfoo"、インデックス 9 に "foobar" という、単語 "foo" と "bar" をそれぞれ1回ずつ使った連結部分文字列が存在することを示しています。

計算量について

各ウィンドウ位置ごとに ok() 関数が O(m × k)(m は単語数、k は1単語の長さ)の処理を行うため、全体の時間計算量は O(n × m × k) となります(n は文字列 s の長さ)。さらに最適化したい場合は、単語単位でウィンドウをスライドさせる手法を採用することで、O(n × k) まで計算量を抑えることが可能です。