max element in specific structure

I have array of length n, from which I build such sequence b that:

b[0] = 0
b[1] = b[0] + a[0]
b[2] = b[1] + a[0]
b[3] = b[2] + a[1]
b[4] = b[3] + a[0]
b[5] = b[4] + a[1]
b[6] = b[5] + a[2]
b[7] = b[6] + a[0]
b[8] = b[7] + a[1]
b[9] = b[8] + a[2]
b[10] = b[9] + a[3]
#etc.

a can contain non-positive values. I need to find max element of b. I only came up with O(n^2) solution. Is there faster approach?

def find_max(a):
  b = [0]
  i = 0
  count = 0
  while i < len(a):
    j = 0
    while j <= i:
      b.append(b[count] + a[j])
      count += 1
      j += 1
    i += 1
  return max(b)

Solution

O(n) time and O(1) space.

Consider this (outer loop) round:

b[4] = b[3] + a[0]   = b[3] + a[0]
b[5] = b[4] + a[1]   = b[3] + a[0] + a[1]
b[6] = b[5] + a[2]   = b[3] + a[0] + a[1] + a[2]

You don't need all of these. It's enough to know:

The maximum of them. Which is b[3] + max(prefix sums of a[:3]).
The last of them, b[6] = b[3] + sum(a[:3]). Because you need that for the next round.

In general, to find the maximum of each round, it's enough to know:

The b value which the round starts with.
The max prefix sum of a[:...].

Add them together to know the max b-value in the round. And return the maximum of these rounds maximums.

We can update these values in O(1) time for each round:

def find_max_linear(a):
  b = max_b = 0 
  sum_a = max_sum_a = 0
  for x in a:
    sum_a += x
    max_sum_a = max(max_sum_a, sum_a)
    max_b = max(max_b, b + max_sum_a)
    b += sum_a
  return max_b

Testing:

import random
for _ in range(10):
  a = random.choices(range(-100, 101), k=100)
  expect = find_max(a)
  result = find_max_linear(a)
  print(result == expect, expect, result)

Output (Attempt This Online!):

True 8277 8277
True 2285 2285
True 5061 5061
True 19261 19261
True 0 0
True 0 0
True 47045 47045
True 531 531
True 703 703
True 24073 24073

Fun oneliner (also O(n) time, but O(n) space due to the unpacking):

from itertools import accumulate as acc
from operator import add

def find_max_linear(a):
  return max(0, *map(add, acc(acc(a), initial=0), acc(acc(a), max)))

Or broken into a few lines with comments:

def find_max_linear(a):
  return max(0, *map(
    add,
    acc(acc(a), initial=0),  # each round's initial b-value
    acc(acc(a), max)         # each round's max prefix sum of a[:...]
  ))

Attempt This Online!