特定结构中的最大元素

Question

我有长度为 n 的数组，从中构建这样的序列 b：

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 可以包含非正值。我需要找到 b 的最大元素。我只想出了 O(n^2) 的解决方案。有没有更快的方法？

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)

Answer 1

O(n) 时间和 O(1) 空间。

考虑这一轮（外循环）：

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]

你不需要所有这些。知道就足够了：

• 其中最大的。这是b[3] + max(prefix sums of a[:3]).
• 他们中的最后一个，b[6] = b[3] + sum(a[:3])。因为你下一轮需要它。

一般来说，要找到每轮的最大值，只需知道：

• b该回合开始的值。
• 的最大前缀和a[:...]。

将它们加在一起即可知道b本轮中的最大值。并返回这些轮数最大值中的最大值。

我们可以在每一轮的 O(1) 时间内更新这些值：

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

测试：

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)

输出（在线尝试！）：

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

有趣的 oneliner（也是 O(n) 时间，但由于解包而占用 O(n) 空间）：

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]

或者用注释分成几行：

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[:...]
  ))

在线尝试这个！