求数组元素与两整数最小绝对差之和的最优O(N)/O(NlogN)解法

问题描述 #

给定一个包含N个整数的数组A，需找出两个整数x和y，使得数组中每个元素与这两个选定整数之一的绝对差之和最小，即计算表达式Σ(min(|a[i] - x|, |a[i] - y|))的最小值。

示例 #

• 示例1：
  • N = 4
  • A = [2,3,6,7]
  • 最优解：选择3和7，总和为|2-3| + |3-3| + |6-7| + |7-7| = 2
• 示例2：
  • N = 8
  • A = [2, 3, 5, 8, 11, 14, 17, 996]
  • 最优解：选择8和996，总和为6+5+3+0+3+6+9+0 = 32

约束条件 #

• 1<=T<=100
• 2<=N<=5*10^3
• 1<=A[i]<=10^5
• 所有测试用例的N之和不超过5*10^3

原代码分析 #

你提供的代码采用三重循环，时间复杂度为O(N³)。当N达到5000时，运算量会达到1250亿次，完全无法在合理时间内运行，必须进行优化。

优化思路 #

1. 排序数组：排序后，最优的分组方式一定是将数组分成连续的左右两部分。左半部分选一个点x，右半部分选一个点y——若a<=b<=c且x<=y，若b选x更优则a必选x，若b选y更优则c必选y，分组必然连续。
2. 前缀和预处理：利用前缀和数组快速计算任意区间的元素和，避免重复计算绝对差和。
3. 枚举分割点+中位数优化：对于每个分割点，左半部分的最优x是左半部分的中位数（中位数能使区间内元素的绝对差和最小），右半部分的最优y是右半部分的中位数。枚举所有可能的分割点，计算对应总和并取最小值即可。

优化后的代码 #

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

using namespace std;

// 计算区间[l, r]（闭区间）内，选中位数的绝对差和
long long calculateSum(const vector<int>& arr, const vector<long long>& prefix, int l, int r) {
    int len = r - l + 1;
    int mid = l + len / 2; // 取中间位置元素作为中位数（偶数长度取靠右的，不影响结果）
    // 左边元素的差和：中位数*左边元素个数 - 左边元素和
    long long leftSum = (long long)arr[mid] * (mid - l + 1) - (prefix[mid + 1] - prefix[l]);
    // 右边元素的差和：右边元素和 - 中位数*右边元素个数
    long long rightSum = (prefix[r + 1] - prefix[mid + 1]) - (long long)arr[mid] * (r - mid);
    return leftSum + rightSum;
}

int minAbsDiff(vector<int> Arr, int N) {
    sort(Arr.begin(), Arr.end());
    // 前缀和数组：prefix[0]=0，prefix[i] = Arr[0]+...+Arr[i-1]
    vector<long long> prefix(N + 1, 0);
    for (int i = 0; i < N; ++i) {
        prefix[i + 1] = prefix[i] + Arr[i];
    }
    
    long long minTotal = LLONG_MAX;
    // 枚举分割点：左半部分[0, k]，右半部分[k+1, N-1]，k从0到N-2
    for (int k = 0; k < N - 1; ++k) {
        long long leftSum = calculateSum(Arr, prefix, 0, k);
        long long rightSum = calculateSum(Arr, prefix, k + 1, N - 1);
        long long total = leftSum + rightSum;
        if (total < minTotal) {
            minTotal = total;
        }
    }
    
    cout << "Sum is :" << minTotal << endl;
    return (int)minTotal;
}

// 测试示例
int main() {
    vector<int> arr1 = {2,3,6,7};
    minAbsDiff(arr1, 4); // 输出2
    
    vector<int> arr2 = {2,3,5,8,11,14,17,996};
    minAbsDiff(arr2, 8); // 输出32
    return 0;
}

代码解释 #

1. 排序：对数组排序，时间复杂度O(NlogN)，是算法的主要耗时部分。
2. 前缀和数组：O(N)时间预处理，快速计算任意区间的元素和。
3. calculateSum函数：O(1)时间计算指定区间选中位数的绝对差和，利用前缀和避免遍历元素。
4. 枚举分割点：O(N)时间遍历所有可能的分组方式，计算总绝对差和并取最小值。

整个算法的时间复杂度为O(NlogN)，完全满足题目约束，运行效率比原代码提升了几个数量级。

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