如何基于正/负变化将numpy数组拆分为2d数组

k3fezbri  于 2023-10-19  发布在  其他
关注(0)|答案(3)|浏览(107)

我有一个numpy 1D数组:

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]]
xbp102n0

xbp102n01#

你可以使用array_splitdiffnonzero

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

Ouptut:

[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]]

方法比较

mccptt67

mccptt672#

如果你想要标准的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]]
igetnqfo

igetnqfo3#

你可以尝试以下逻辑:

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

相关问题