Number of subsequences of length 3 or more, such that the sum of its "middle" is equal to the values of its "endpoints"

Question

Given an array of integers (length N), return the count of subsequences of length at least 3 (from 3 to N) such that the following conditions hold: the value of the first element of the subsequence equals the value of the last element of the subsequence, and they both equal the sum of the values of the 'middle' elements of the subsequence, i.e., excluding the first and last elements.

You are not allowed to sort or change the sequence of the array of integers in any way.

What would be a more effective way to solve the problem than the following "naive" approach?

def count_balanced_subsegments(capacity):
    n = len(capacity)
    count = 0

    prefix_sum = [0] * n
    prefix_sum[0] = capacity[0]

    # Calculate the prefix sum of the capacity array
    for i in range(1, n):
        prefix_sum[i] = prefix_sum[i - 1] + capacity[i]

    for i in range(n - 2):  # Start from the first server
        for j in range(i + 2, n):  # Try all possible subsegment lengths starting from 3
            if capacity[i] == capacity[j] and prefix_sum[j - 1] - prefix_sum[i] == capacity[i]:
                count += 1

    return count

# Example usage:
capacity = [9, 3, 3, 3, 9]
result = count_balanced_subsegments(capacity)
print(result)  # Output should be 2

Answer 1

One potential to speed-up the computation is to compute the combinations between indices of equal value.

from itertools import combinations

def count_balanced_subsegments_2(capacity):
    indices, prefix_sum = {}, []

    cnt = 0
    for i, v in enumerate(capacity):
        indices.setdefault(v, []).append(i)
        prefix_sum.append(cnt := cnt + capacity[i])

    count = 0
    for k, v in indices.items():
        for i, j in combinations(v, 2):
            if j - i >= 2 and prefix_sum[j - 1] - prefix_sum[i] == capacity[i]:
                count += 1

    return count

---

Quick benchmark, with N = 50_000 and values between 1 and 10:

from random import randint, seed
from timeit import timeit

seed(42)
capacity = [randint(1, 10) for _ in range(50_000)]

assert count_balanced_subsegments(capacity) == count_balanced_subsegments_2(capacity)

t1 = timeit("count_balanced_subsegments(capacity)", number=1, globals=globals())
t2 = timeit("count_balanced_subsegments_2(capacity)", number=1, globals=globals())

print(f"{t1=}")
print(f"{t2=}")

Prints:

t1=28.380160036991583
t2=5.308306446007919

Answer 2

We can store prefix sums in dictionary together with item value as tuple key(psum, a[i]). Value for this key is the number of correponding prefix sums.

For every new value a[i] we check that difference psum - a[i] already is stored in dictionary (analog of yours prefix_sum[j - 1] - prefix_sum[i] == capacity[i]) with border value equal to a[i], and update counter if such sums do exist.

Delay of dictionary update (last tuple storing) is intended to provide subsize>=3 condition

Time is linear, ~50 milliseconds for list size 100 000

from collections import defaultdict
def cnt_balanced(a):
    n = len(a)
    count = 0
    psum = a[0]
    last = (psum,a[0])
    sums = defaultdict(int)

    for i in range(1, n):
        t = (psum - a[i], a[i])
        if t in sums:
            count += sums[t]

        sums[last] += 1
        psum += a[i]
        last = (psum, a[i])

    return count