Count the valid segments for input array of numbers

Question

I have an array of numbers of size n, I want to count the segments, (which are subarrays of size 3 or more) based on below condition:

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

it means for selected sub-array arr[left, right] the values of endpoints (left, and right) should be same and this value should match with sum of remaining items in subarray.

Example:

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

Valid segments:

[3,3,3] 

-> here sub-arr[left, right], arr[left] = 3, arr[right] 3, sum(arr[left+1, right-1]) = 3

[9,3,3,3,9]

-> here sub-arr[left, right], arr[left] = 9, arr[right] 9, sum(arr[left+1, right-1]) = 3+3+3 = 9

Result = 2

Constraints: 1 <= n <= 3 * 10^5 1 <= arr[i] <= 10^9

Here is my code that takes time of O(n^2)

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;
    }
}

How can I solve this in less time complexity.

Answer 1

We can try every index as the possible right index of a segment. Let prefSum[i] indicate the sum of all elements from index 0 up to index i in the array. The sum of the subarray from index l to r is prefSum[r] - prefSum[l - 1]. For any left index to form a valid segment with a given right index, we require prefSum[r - 1] - prefSum[l] (subarray sum from l + 1 to r - 1) to equal arr[l]. Since arr[l] = arr[r] is also required, we can solve for prefSum[l] and find that it should be prefSum[r - 1] - arr[r]. Now, we can store the counts of pairs of arr[i], prefSum[i] for each index in a HashMap. Overall, this yields a linear solution.

public static int solve(int[] arr) {
    if (arr == null || arr.length < 3) return 0;
    record Value(int elem, int prefSum) {}
    var counts = new HashMap<Value, Integer>();
    int sum = arr[0], result = 0;
    for (int i = 1; i < arr.length; i++) {
        result += counts.getOrDefault(new Value(arr[i], sum - arr[i]), 0);
        counts.merge(new Value(arr[i - 1], sum), 1, Integer::sum);
        sum += arr[i];
    }
    return result;
}