ys memos

Blog

Leetcode 15 3sum


leetcode

2020/05/27


vector<int> numsが与えられ、その内3つの要素の和が0になる組を求める問題

numsの要素は重複することがあるが、同じ組み合わせは解に含めない。


vector<vector<int>> threeSum(vector<int>& nums) {
  if(nums.size() < 3) return {};
  map<int,int> mp;
  for(const auto& num : nums) mp[num]++;
  vector<vector<int>> threes;
  for(auto itr_i=mp.begin(); itr_i!=mp.end(); ++itr_i){
    itr_i->second--;
    for(auto itr_j=itr_i; itr_j!=mp.end(); ++itr_j){
      if( itr_j->second == 0 ) continue;
      itr_j->second--;
      int nums_k = - itr_i->first - itr_j->first;
      if( nums_k < itr_j->first ){
        itr_j->second++;
        break;
      }
      if( mp[nums_k] != 0 ){
        threes.push_back({itr_i->first, itr_j->first, nums_k});
      }
      itr_j->second++;
    }
    itr_i->second++;
  }
  return threes;
}

numsの要素数が3未満であれば組み合わせを作ることができないので空の配列を返す。

vector<vector<int>> threeSum(vector<int>& nums) {
  if(nums.size() < 3) return {};

std::mapに、数値の個数をカウントしていく。

  map<int,int> mp;
  for(const auto& num : nums) mp[num]++;

ここからが組み合わせの探索。 3つの数を{a, b, c}とした時、 $$a \leq b \leq c$$ を満たすものとする。この順番を維持することで、組み合わせの重複を防ぐことができる。

std::mapのイテレータはキーの小さい方から順に走査してくれる。 そのキーの値を、「選択中」という意味で一つマイナスしておく(同値が含まれる組み合わせのため)。

  vector<vector<int>> threes;
  for(auto itr_i=mp.begin(); itr_i!=mp.end(); ++itr_i){
    itr_i->second--;

bは、aと同じ値になる可能性もあるので、itr_jitr_iから走査する。 ここでも、「選択中」という意味で一つマイナスしておく。

    for(auto itr_j=itr_i; itr_j!=mp.end(); ++itr_j){
      if( itr_j->second == 0 ) continue;
      itr_j->second--;

選択中のabに対して和が0となるcを計算する(普通の引き算0-a-b)。 ここで、

$$ a \leq b \leq c $$

の関係が崩れた場合は、それ以上bを走査しても意味がないのでbreakして次のaに対して走査する。

      int nums_k = - itr_i->first - itr_j->first;
      if( nums_k < itr_j->first ){
        itr_j->second++;
        break;
      }

abcの関係が成り立った時、そのcが元の配列に存在している場合は、その組み合わせをthreesに追加する。

      if( mp[nums_k] != 0 ){
        threes.push_back({itr_i->first, itr_j->first, nums_k});
      }

itr_iitr_jを次の値にする。

      itr_j->second++;
    }
    itr_i->second++;
  }

全ての組み合わせthreesを返す。

  return threes;
}

関連タグを探す