优化有效元素对算法时间复杂度，解决超时问题

给定整数n、m、k和一个长度为n的整数数组arr，当元素对(arr[i], arr[j])满足以下两个条件时视为有效对：

1. abs(arr[i] - arr[j]) <= m
2. arr[i] + arr[j] <= k
   需要输出有效对的总数量。

示例 #

输入：

4 1 5
1 3 2 3

输出：3
有效对为(1,2)、(3,2)、(2,3)。

现有代码问题 #

当前代码可正常运行，但时间复杂度为O(N²)，会触发超时错误，代码如下：

#include <iostream>
#include <vector>
#include <string>
#include <unordered_map>
#include <algorithm>

using namespace std;


int main() {
    int n, m, k;
    cin >> n >> m >> k;

    vector<int> arr(n);
    for (int i = 0; i < n; ++i) {
        cin >> arr[i];
    }
    sort(begin(arr), end(arr));

    int res = 0;
    for (int i = 0; i < n; ++i) {
        for (int j = i + 1; j < n; ++j) {
            if (arr[i] + arr[j] <= k and abs(arr[i] - arr[j]) <= m) {
                res++;
            } else {
                break;
            }
        }
    }
    cout << res;
    return 0;
}

优化需求 #

希望借助排序思想将时间复杂度降至O(NlogN)，该如何实现？

你已经对数组做了排序，这是优化的核心基础。接下来可以用二分查找替代内层循环，将时间复杂度降到O(NlogN)。

核心思路 #

数组排序后，对于每个元素arr[i]，我们只需找到所有j > i且满足以下两个条件的arr[j]：

1. 因为数组升序，arr[j] >= arr[i]，所以abs(arr[i]-arr[j]) = arr[j]-arr[i] <= m → arr[j] <= arr[i] + m
2. arr[i] + arr[j] <= k → arr[j] <= k - arr[i]

综合两个条件，arr[j]的上限是min(arr[i]+m, k-arr[i])。由于数组有序，我们可以用二分查找快速定位到第一个大于该上限值的元素位置，这个位置之前的所有j > i的元素都满足条件，数量就是该位置与i+1的差值。累加所有i对应的数量即可得到答案。

优化后代码 #

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

int main() {
    int n, m, k;
    cin >> n >> m >> k;

    vector<int> arr(n);
    for (int i = 0; i < n; ++i) {
        cin >> arr[i];
    }
    sort(arr.begin(), arr.end());

    long long res = 0;  // 避免大数量级时int溢出
    for (int i = 0; i < n; ++i) {
        int upper_limit = min(arr[i] + m, k - arr[i]);
        // 找到第一个大于upper_limit的元素迭代器
        auto it = upper_bound(arr.begin() + i + 1, arr.end(), upper_limit);
        // 累加当前i对应的有效对数量
        res += (it - (arr.begin() + i + 1));
    }
    cout << res;
    return 0;
}

复杂度说明 #

• 排序阶段：O(NlogN)
• 遍历+二分查找阶段：每个二分查找耗时O(logN)，共N次，总耗时O(NlogN)
• 整体时间复杂度：O(NlogN)
• 空间复杂度：O(1)（忽略排序的栈空间开销）

关键注意点 #

• 使用long long存储结果：当n达到1e5级别时，有效对数量可能超过int的最大值（约2e9），必须用64位整数避免溢出。
• upper_bound的使用：该函数返回第一个大于目标值的迭代器，因此有效元素的数量是迭代器与起始位置（arr.begin()+i+1）的差值，无需额外调整。

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