如何根据正负变化将一个numpy数组拆分为二维数组

Question

我有一个numpy的一维数组：

import numpy as np

arr = np.array([1,  1,  3, -2, -1,  2,  0,  2,  1,  1, -3, -1,  2])

我想根据数组元素正负值的变化（0位于正值范围内）将其拆分为另一个二维数组。但是，原始元素的顺序应该保持不变。

期望的结果是：

new_arr = [[1, 1, 3], [-2, -1], [2, 0, 2, 1, 1], [-3, -1], [2]]

- Sun Jar

Answer 1

你可以使用 array_split, diff, nonzero 这几个函数。

np.array_split(arr, np.nonzero(np.diff(arr>=0))[0]+1)

输出：

[array([1, 1, 3]),
 array([-2, -1]),
 array([2, 0, 2, 1, 1]),
 array([-3, -1]),
 array([2])]

中级：

# arr>0
[ True  True  True False False  True False  True  True  True False False  True]

# np.diff(arr>=0)
[False False  True False  True False False False False  True False  True]

# np.nonzero(np.diff(arr>=0))[0]+1
[ 3  5 10 12]

对于输出的列表：

out = list(map(list, np.array_split(arr, np.nonzero(np.diff(arr>=0))[0]+1)))

输出：

[[1, 1, 3], [-2, -1], [2, 0, 2, 1, 1], [-3, -1], [2]]

或者使用itertools.groupby：

from itertools import groupby

out = [list(g) for _,g in groupby(arr, key=lambda x: x>=0)]

输出：

[[1, 1], [3, -2], [-1, 2, 0, 2, 1], [1, -3], [-1, 2]]

方法的比较

- mozway

---

可以从新的结果中删除那些“数组”吗？ - undefined

@Sun 你是指这个代码吗？out = list(map(list, np.array_split(arr, np.nonzero(np.diff(arr>=0))[0]+1))) - undefined

但是如果你想要一个列表，你也可以使用itertools.groupby。 - undefined

这就是我想要的，我更喜欢np方法而不是itertools。 - undefined

@Sun 是的，对于大型数组来说速度会稍快一些。我会添加一个比较。 - undefined

Answer 2

如果你想要标准的Python列表，你可以使用itertools.groupby。

from itertools import groupby

arr = np.array([1, 1, 3, -2, -1, 2, 0, 2, 1, 1, -3, -1, 2])

out = []
for _, g in groupby(arr, lambda k: k < 0):
    out.append(list(g))

print(out)

输出：

[[1, 1, 3], [-2, -1], [2, 0, 2, 1, 1], [-3, -1], [2]]

- Andrej Kesely

---

我刚刚添加了相同的内容。如果需要使用Python列表，我同意这可能会更直接。 - undefined

还是一个不错的算法，但我更喜欢np原生方法，itertools可能有点慢。 - undefined

Answer 3

也许你可以尝试以下的逻辑：

def split_array(arr):
    new_arr = []
    subarr = []
    
    for elem in arr:
        if (elem >= 0 and subarr and subarr[-1] < 0) or (elem < 0 and subarr and subarr[-1] >= 0):
            new_arr.append(subarr)
            subarr = []
        subarr.append(elem)
    
    if subarr:
        new_arr.append(subarr)
    
    return new_arr

- Giba