如何将满足端点值等于中间和的数组有效子数组计数算法优化至更优复杂度？

问题描述 #

给定长度为n的数值数组，统计满足以下条件的有效段（即长度≥3的子数组）：

arr[left] = arr[right] = sumOfItems[left+1, right-1]

即选中的子数组arr[left, right]的两端元素值相等，且该值等于区间left+1到right-1内所有元素的和。

示例 #

输入数组：

arr = [9,3,3,3,9]

有效段：

1. [3,3,3]：arr[left]=3，arr[right]=3，中间元素和为3，满足条件
2. [9,3,3,3,9]：arr[left]=9，arr[right]=9，中间元素和为3+3+3=9，满足条件

统计结果：2

约束条件 #

• 1 ≤ n ≤ 3 * 10^5
• 1 ≤ arr[i] ≤ 10^9

原O(n²)解法的问题 #

原Java代码通过两层循环枚举所有可能的start和end，再计算中间元素和，时间复杂度为O(n²)，对于n=3e5的规模会直接超时，必须优化。

原代码：

public class Main {

    public static void main(String[] args) {
        int[] arr = {9,3,3,3,9};

        int result = solve(arr);
        System.out.println(result);
    }

    public static int solve(int[] arr) {
        int result = 0;

        for (int start = 0; start < arr.length; start++) {
            for (int end = start+2; end < arr.length; end++) {
                if (arr[start] == arr[end]) {
                    long sum = 0;
                    for(int i=start+1; i<end; i++) {
                        sum += arr[i];
                        if(sum > arr[start]) break;
                    }
                    if(sum == arr[start]) {
                        result++;
                    }
                }
            }
        }

        return result;
    }
}

优化思路（O(n log n)时间复杂度） #

利用数组元素全为正整数的特性，前缀和数组是严格递增的，结合哈希表分组+二分查找实现高效统计：

1. 前缀和转换：计算前缀和数组prefix，其中prefix[0] = 0，prefix[k] = arr[0] + arr[1] + ... + arr[k-1]。则区间left+1到right-1的和可表示为prefix[right] - prefix[left+1]。
2. 条件变形：根据题目条件，arr[left] = arr[right]且arr[left] = prefix[right] - prefix[left+1]，代入arr[left] = prefix[left+1] - prefix[left]，可推导出：

   prefix[right] = 2 * prefix[left+1] - prefix[left]
   

3. 哈希表分组：用哈希表存储每个元素值对应的所有索引及对应前缀和（prefix[right+1]），因为相同值的元素才可能成为left和right的配对。
4. 二分查找快速定位：对于每个left，计算目标前缀和target_prefix，再在哈希表中对应值的列表里，用二分查找找到所有满足right > left+1且prefix[right+1] == target_prefix的right数量，累加至结果。

优化后的Java代码 #

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

public class OptimizedSolution {

    public static void main(String[] args) {
        int[] arr = {9, 3, 3, 3, 9};
        System.out.println(countValidSegments(arr));
    }

    public static int countValidSegments(int[] arr) {
        int n = arr.length;
        long[] prefix = new long[n + 1];
        // 计算前缀和
        for (int i = 0; i < n; i++) {
            prefix[i + 1] = prefix[i] + arr[i];
        }

        // 哈希表：key是数组元素值，value是存储(索引i, prefix[i+1])的列表
        Map<Long, List<long[]>> valueMap = new HashMap<>();
        for (int i = 0; i < n; i++) {
            long val = arr[i];
            valueMap.computeIfAbsent(val, k -> new ArrayList<>()).add(new long[]{i, prefix[i + 1]});
        }

        int result = 0;
        for (int left = 0; left < n; left++) {
            long targetVal = arr[left];
            // 计算目标前缀和
            long targetPrefix = 2 * prefix[left + 1] - prefix[left];
            List<long[]> candidates = valueMap.get(targetVal);
            if (candidates == null) {
                continue;
            }

            // 二分查找找到第一个right > left+1的位置
            int low = 0, high = candidates.size();
            while (low < high) {
                int mid = (low + high) / 2;
                if (candidates.get(mid)[0] > left + 1) {
                    high = mid;
                } else {
                    low = mid + 1;
                }
            }

            // 在low到末尾的范围内，找prefix[right+1] == targetPrefix的数量
            // 因为前缀和严格递增，相同targetPrefix的元素是连续的
            int leftBound = low;
            low = leftBound;
            high = candidates.size();
            while (low < high) {
                int mid = (low + high) / 2;
                if (candidates.get(mid)[1] >= targetPrefix) {
                    high = mid;
                } else {
                    low = mid + 1;
                }
            }
            int start = low;

            low = start;
            high = candidates.size();
            while (low < high) {
                int mid = (low + high) / 2;
                if (candidates.get(mid)[1] > targetPrefix) {
                    high = mid;
                } else {
                    low = mid + 1;
                }
            }
            int end = low;

            result += (end - start);
        }

        return result;
    }
}

代码说明 #

• 前缀和用long类型存储，避免因元素值过大导致溢出。
• 哈希表分组存储相同值的元素信息，减少无效遍历。
• 两次二分查找分别定位符合条件的right的起始范围和目标前缀和的边界，利用前缀和严格递增的特性，确保查找效率为O(log m)（m为相同值元素的数量）。
• 整体时间复杂度为O(n log n)，可以轻松处理n=3e5的规模。

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