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_j
はitr_i
から走査する。
ここでも、「選択中」という意味で一つマイナスしておく。
for(auto itr_j=itr_i; itr_j!=mp.end(); ++itr_j){
if( itr_j->second == 0 ) continue;
itr_j->second--;
選択中のa
とb
に対して和が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;
}
a
・b
・c
の関係が成り立った時、そのc
が元の配列に存在している場合は、その組み合わせをthrees
に追加する。
if( mp[nums_k] != 0 ){
threes.push_back({itr_i->first, itr_j->first, nums_k});
}
itr_i
とitr_j
を次の値にする。
itr_j->second++;
}
itr_i->second++;
}
全ての組み合わせthrees
を返す。
return threes;
}