将一个整数列表缩短为连续正数或负数之和。

Question

我想编写一个处理整数列表的函数，最好的方式是通过示例展示：

input [0,1,2,3, -1,-2,-3, 0,1,2,3, -1,-2,-3] will return [6,-6,6,-6]

我这里有一份可以真正起作用的草稿：

def group_pos_neg_list(nums):
    p_nums = []

    # to determine if the first element >=0 or <0
    # create pos_combined and neg_combined as a list to check the length in the future
    if nums[0] >= 0:
        pos_combined, neg_combined = [nums[0]], []
    elif nums[0] < 0:
        pos_combined, neg_combined = [], [nums[0]]

    # loop over each element from position 1 to the end
    # accumulate pos num and neg nums and set back to 0 if next element is different
    index = 1
    while index < len(nums):
        if nums[index] >= 0 and nums[index-1] >= 0: # both posivite
            pos_combined.append(nums[index])
            index += 1
        elif nums[index] < 0 and nums[index-1] < 0: # both negative
            neg_combined.append(nums[index])
            index += 1
        else:
            if len(pos_combined) > 0:
                p_nums.append(sum(pos_combined))
                pos_combined, neg_combined = [], [nums[index]]
            elif len(neg_combined) > 0:
                p_nums.append(sum(neg_combined))
                pos_combined, neg_combined = [nums[index]], []
            index += 1

    # finish the last combined group
    if len(pos_combined) > 0:
        p_nums.append(sum(pos_combined))
    elif len(neg_combined) > 0:
        p_nums.append(sum(neg_combined))

    return p_nums

但我对此并不太满意，因为它看起来有点复杂。特别是代码中有一部分是重复的：

if len(pos_combined) > 0:
    p_nums.append(sum(pos_combined))
    pos_combined, neg_combined = [], [nums[index]]
elif len(neg_combined) > 0:
    p_nums.append(sum(neg_combined))
    pos_combined, neg_combined = [nums[index]], []

我需要写两次，因为循环不会计算最后一组整数，所以需要额外的步骤。

有没有简化这个过程的方法？

- jxie0755

---

可以使用 itertools.groupby 轻松完成此操作。只需使用列表推导式将正数和负数的运行总和相加即可。 - Christian Dean

@ChristianDean 我确实知道如何使用groupby，但我正在尝试不依赖现成的模块自己写出来。 - jxie0755

你的原始列表已排序保证吗？换句话说，它是否具有结构 [1, 2, 3, -1, -2, -3, ...]？ - Christian Dean

@ChristianDean 不，没有排序。完全随机。 - jxie0755

Answer 1

使用groupby

不需要那么复杂：我们可以先按符号进行groupby，然后再计算总和，所以：

from itertools import groupby

[sum(g) for _, g in groupby(data, lambda x: x >= 0)]

这将产生以下结果：

>>> from itertools import groupby
>>> data = [0,1,2,3, -1,-2,-3, 0,1,2,3, -1,-2,-3]
>>> [sum(g) for _, g in groupby(data, lambda x: x >= 0)]
[6, -6, 6, -6]

所以，groupby生成带有“键”（我们使用lambda计算的部分）和“burst”的可迭代元组（具有相同键的连续子序列）。我们只对后者 g 感兴趣，然后计算 sum(g) 并将其添加到列表中。

自定义算法

我们也可以编写自己的版本，方法如下：

swap_idx = [0]
swap_idx += [i+1 for i, (v1, v2) in enumerate(zip(data, data[1:]))
             if (v1 >= 0) != (v2 >= 0)]
swap_idx.append(None)

our_sums = [sum(data[i:j]) for i, j in zip(swap_idx, swap_idx[1:])]

在这里，我们首先构建一个列表 swap_idx，其中存储符号更改的元素的索引。因此，对于您的示例代码，它是:

>>> swap_idx
[0, 4, 7, 11, None]

0和None是代码明确添加的。因此，既然我们已经确定了符号变化的点，我们可以将这些子序列相加，使用sum(data[i:j])。因此，我们使用zip(swap_idx, swap_idx[1:])来获取两个连续的索引，然后我们可以将该切片相加。

更详细的版本

以上内容不太易读：虽然它能工作，但需要一些推理。我们也可以生成更详细的版本，并使其更加通用，例如：

def groupby_aggregate(iterable, key=lambda x: x, aggregate=list):
    itr = iter(iterable)
    nx = next(itr)
    kx = kxcur = key(nx)
    current = [nx]
    try:
        while True:
            nx = next(itr)
            kx = key(nx)
            if kx != kxcur:
                yield aggregate(current)
                current = [nx]
                kxcur = kx
            else:
                current.append(nx)
    except StopIteration:
         yield aggregate(current)

我们可以这样使用它：

list(groupby_aggregate(data, lambda x: x >= 0, sum))

- Willem Van Onsem

---

@Code_Control_jxie0755：那么你完全违背了Python的精神。 - Willem Van Onsem

哦，别这样@Willem :-) 我认为OP只是想练习自己创建函数。如果他正在编写实际的生产代码，那么我会同意你的观点。 - Christian Dean

我理解精神，只是我在这里尝试学习算法... - jxie0755

在这种情况下，您可以检查groupby的工作原理。该文档中提供了与代码等效的显式逻辑。 - Moinuddin Quadri

@Code_Control_jxie0755：itertools是一个标准库，通常情况下应该在文件顶部导入，除非你引入了循环导入，或者有（非常）好的理由不立即导入该模块。 - Willem Van Onsem

Answer 2

您可以使用itertools.groupby，利用键来按所有大于或等于零的值进行分组：

import itertools
s = [0,1,2,3, -1,-2,-3, 0,1,2,3, -1,-2,-3]
new_s = [sum(b) for a, b in itertools.groupby(s, key=lambda x: x >=0)]

输出：

[6, -6, 6, -6]

- Ajax1234

---

谢谢！我知道如何使用groupby，但如果我想编写自己的版本呢？ - jxie0755

@ChristianDean 谢谢你指出这个问题。请查看我的最近编辑。 - Ajax1234

Answer 3

以下是一种不需要任何外部导入，仅使用 reduce() 的方法：

def same_sign(a, b):
    """Returns True if a and b have the same sign"""
    return (a*b>0) or (a>=0 and b>=0)

l = [0,1,2,3, -1,-2,-3, 0,1,2,3, -1,-2,-3] 
reduce(
    lambda x, y: (x+y if same_sign(x,y) else [x, y]) if not isinstance(x, list) else x[:-1] + [x[-1] + y] if same_sign(x[-1],y) else x + [y],
    l
)
#[6, -6, 6, -6]

解释

这有点难以解释，但我会尝试。

从文档中调用reduce()将会:

从左到右将可迭代对象的项累积地应用于两个参数的函数

在这种情况下，我从您的列表中取出两个值(x和y)，并执行以下操作:

• 如果x不是一个list:
  • 如果x和y具有相同的符号(乘积 >=0)，则将它们相加
  • 否则返回一个列表[x, y]
• 如果x是一个list，则仅修改x的最后一个元素。
  • 如果符号匹配，则添加y。
  • 否则向列表x附加一个新元素

注意

您可能 不应该 这样做，因为代码很难阅读和理解。我只是想表明这是可能的。

---

更新

同样代码的更易读版本:

def same_sign(a, b):
    """Returns True if a and b have the same sign"""
    return (a*b>0) or (a>=0 and b>=0)

l = [0,1,2,3, -1,-2,-3, 0,1,2,3, -1,-2,-3] 
def reducer(x, y):
    if isinstance(x, list):
        if same_sign(x[-1], y):
            return x[:-1] + [x[-1] + y]
        else:
            return x + [y]
    else:
        if same_sign(x, y):
            return x+y
        else:
            return [x, y]
reduce(reducer, l)
#[6, -6, 6, -6]

- pault

---

谢谢，还不错！只是也许你可以将它写成多行……如果一行看起来太隐晦的话，没必要把所有东西都写在一行里。 - jxie0755

这就是为什么在我的代码中，我使用了if和elif条件来确保两个相邻的元素都是正数或负数，否则进行其他处理。 - jxie0755

你可以检查我的原始代码，它们可以工作，只是不太优雅 :(，但仍然是O(n)算法。 - jxie0755

@Code_Control_jxie0755 对于 [0,0,-3,1] 它得到了正确的输出。更好的情况是 [1, 1, -2, -2, 0,0,-3,1]，它输出了 [2, -7, 1] 而不是 [2, -4, 0, -3, 1]（我认为这是你想要的）。我会进行更新。 - pault

@Code_Control_jxie0755 我进行了一次编辑，应该可以处理这个边缘情况。或者你可以使用 np.sign()，但那需要一个导入。 - pault