如何降低计算数组元素绝对差之和的代码时间复杂度？

给定整数数组A和B，需为B中的每个元素计算其与A中所有元素的绝对差之和，构建结果数组。例如：

• A = [1,2,3]，B = [3,2,1,5]
• 结果数组为[3,2,3,9]，计算逻辑：
  • 3与A的绝对差之和：|3-1|+|3-2|+|3-3|=2+1+0=3
  • 2与A的绝对差之和：|2-1|+|2-2|+|2-3|=1+0+1=2
  • 1与A的绝对差之和：|1-1|+|1-2|+|1-3|=0+1+2=3
  • 5与A的绝对差之和：|5-1|+|5-2|+|5-3|=4+3+2=9

当前使用HashMap避免重复计算的代码如下：

static List<Long> solve(int[] A, int[] B) {
    List<Long> result = new ArrayList<>();
    Map<Integer, Long> map = new HashMap<>();
    for(int b: B) {
        long sum = 0;
        if(map.get(b) != null) {
            result.add(map.get(b));
        } else {
            for(int a : A) {
                sum += Math.abs(b - a);
            }
            map.put(b, sum);
            result.add(sum);
        }
    }
    return result;
}

该代码时间复杂度为O(m*n)（m为数组A的长度，n为数组B的长度），请问如何降低这段代码的时间复杂度？

可以通过排序+前缀和+二分查找的组合将时间复杂度降至O(m log m + n log m)，具体思路如下：

1. 排序数组A：先对A进行升序排序，时间复杂度O(m log m)。排序后可利用有序数组特性，通过二分查找快速定位元素位置，拆分绝对差的计算逻辑。
2. 计算前缀和数组：基于排序后的A生成前缀和数组prefix，其中prefix[i]表示A中前i个元素的累加和（prefix[0] = 0，prefix[1] = A[0]，prefix[2] = A[0]+A[1]，以此类推）。前缀和能快速计算任意区间内元素的总和，避免重复累加。
3. 二分查找计算每个B元素的绝对差和：对于B中的每个元素b：
   • 用二分查找找到A中第一个大于b的元素索引pos，即A中有pos个元素小于等于b，剩余m-pos个元素大于b。
   • 左边（<=b的元素）的绝对差总和：pos * b - prefix[pos]（等价于pos个b的和减去左边元素的总和）。
   • 右边（>b的元素）的绝对差总和：(prefix[m] - prefix[pos]) - (m - pos) * b（等价于右边元素的总和减去(m-pos)个b的和）。
   • 总绝对差和为左右两部分之和。
4. 可选保留HashMap缓存：如果B中存在大量重复元素，依然可以用HashMap缓存已计算过的b的结果，避免重复执行二分查找和计算，进一步提升效率。

优化后的代码示例 #

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class Solution {
    static List<Long> solve(int[] A, int[] B) {
        List<Long> result = new ArrayList<>();
        if (A == null || A.length == 0) {
            for (int b : B) {
                result.add(0L);
            }
            return result;
        }
        
        // 排序数组A
        Arrays.sort(A);
        int m = A.length;
        // 计算前缀和数组
        long[] prefix = new long[m + 1];
        for (int i = 0; i < m; i++) {
            prefix[i + 1] = prefix[i] + A[i];
        }
        
        Map<Integer, Long> cache = new HashMap<>();
        for (int b : B) {
            if (cache.containsKey(b)) {
                result.add(cache.get(b));
                continue;
            }
            // 二分查找定位第一个大于b的元素位置
            int pos = Arrays.binarySearch(A, b);
            if (pos < 0) {
                pos = -pos - 1;
            } else {
                // 处理A中存在重复b的情况，找到最后一个b的下一个位置
                while (pos < m && A[pos] == b) {
                    pos++;
                }
            }
            // 计算左右两部分的绝对差和
            long leftSum = (long) pos * b - prefix[pos];
            long rightSum = (prefix[m] - prefix[pos]) - (long) (m - pos) * b;
            long total = leftSum + rightSum;
            
            cache.put(b, total);
            result.add(total);
        }
        return result;
    }
    
    public static void main(String[] args) {
        int[] A = {1,2,3};
        int[] B = {3,2,1,5};
        System.out.println(solve(A, B)); // 输出 [3, 2, 3, 9]
    }
}

复杂度分析 #

• 排序A的时间：O(m log m)
• 计算前缀和：O(m)
• 处理每个B元素：每个元素的二分查找为O(log m)，加上O(1)的总和计算，n个元素总耗时O(n log m)
• 总时间复杂度：O(m log m + n log m)，相比原O(m*n)的复杂度，在m和n较大时性能提升显著。

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