std::map の範囲 for 文内で operator[] を使わない

バグったのでメモ



突然ですが次のようなコードがあります。

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

int main() {
  map<int, int> mp;
  mp[0] = 1, mp[1] = 2;
  int sum = 0;
  for (auto [key, val] : mp) {
    cout << key << " " << val << "\n";
    sum += mp[key + 1];
  }
  cout << sum << "\n";
}

map の全ての key に対して mp[key+1] を求めるコードです。

mp[0+1] = 2, mp[1+1] = 0 なので sum の値は 3 になりそうですが、実際に実行すると次のようになります。

0 1
1 2
2 0
3 0
4 0
5 0
6 0
7 0
8 0
9 0
10 0
11 0
12 0
13 0
14 0
15 0
16 0
17 0
18 0
19 0
20 0
21 0
...(以下ずっと続く)

std::map の operator[] は「キーに対応する値を返し、対応する値がない場合は要素をデフォルトで構築する」という仕様になっています。そのため mp[key+1] と書いただけで {key+1,0} という要素が構築されてしまい、無限ループになります。
std::unordered_map の場合も同様です。(こちらは収納順序が不定である分バグに気付きにくいので注意)

解決策

find関数やcount関数を使いましょう。findの方が要素の探索が一回で済むので速いと思います。

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

int main() {
  map<int, int> mp;
  mp[0] = 1, mp[1] = 2;
  int sum = 0;
  for (auto [key, val] : mp) {
    cout << key << " " << val << "\n";
    auto itr = mp.find(key + 1);
    if (itr != mp.end()) sum += itr->second;
  }
  cout << sum << "\n";
}
実行結果
0 1
1 2
2

範囲 for 文以外の場面でも無駄に要素が増えると定数倍が大きくなるので、実行時間が気になるときは operator[] を控えましょう。