C++ 使用std::map统计和为指定值的数对的算法如何优化？

• 逻辑错误1：map.insert() 插入重复key会直接失败，重复元素完全无法统计，比如输入两个2、s=4的场景，你的代码会直接丢失重复数据，结果完全错误
• 逻辑错误2：双重遍历的计数逻辑仅在补数大于当前key的场景下能巧合得到正确结果，补数小于当前key的场景会直接漏算；同时无法处理同值元素配对、多重复元素配对的场景
• 性能问题：双重遍历的时间复杂度为O(n²)，当n取上限50000时运算量达到2.5e9级别，会直接超时无法通过测试

我们可以利用map统计元素出现频次，将时间复杂度降到O(n log n)，同时解决上述逻辑问题，核心思路如下：

1. 遍历输入元素时，用map统计每个数字出现的总次数
2. 遍历map的每个key，计算补数target = s - key：
   • 若key < target：计数 += 当前key出现次数 * target出现次数
   • 若key == target：计数 += 组合数cnt * (cnt - 1) / 2，即从cnt个相同元素里选2个的组合数
   • 若key > target：直接跳过，避免重复计数

#include <iostream>
#include <map>
using namespace std;

int main()
{
    cin.tie(0);
    ios::sync_with_stdio(false);
    int s, n;
    cin >> s >> n;
    map<int, int> cnt_map;
    int element;
    for (int i = 0; i < n; ++i) {
        cin >> element;
        cnt_map[element]++; // 重复元素直接累加计数
    }
    long long counter = 0; // 用long long避免大数溢出
    for (auto& p : cnt_map) {
        int key = p.first;
        int cnt = p.second;
        int target = s - key;
        if (cnt_map.count(target) == 0) {
            continue;
        }
        if (key < target) {
            counter += 1LL * cnt * cnt_map[target];
        } else if (key == target) {
            counter += 1LL * cnt * (cnt - 1) / 2;
        }
        // key > target的情况跳过，避免重复统计
    }
    cout << counter << "\n";
    return 0;
}

可选更精简写法 #

也可以采用边插入边查询的写法，省去后续遍历map的步骤，效率更高：

#include <iostream>
#include <map>
using namespace std;

int main()
{
    cin.tie(0);
    ios::sync_with_stdio(false);
    int s, n;
    cin >> s >> n;
    map<int, int> cnt_map;
    long long counter = 0;
    int element;
    for (int i = 0; i < n; ++i) {
        cin >> element;
        int target = s - element;
        // 统计当前元素和之前已插入元素的匹配数，不会重复计数
        if (cnt_map.count(target)) {
            counter += cnt_map[target];
        }
        cnt_map[element]++;
    }
    cout << counter << "\n";
    return 0;
}

内容的提问来源于stack exchange，提问作者MassSpectrometer