python - 根据给定条件拆分 numpy 数组

如何将已排序的 numpy 数组(例如 arr=([5,6,28,29,32,33,87,88,95]))拆分为 sub-arrays,以便始终满足以下两个条件:

(1) 子数组的第一个和最后一个元素之间的差小于 10。

(2) 并且,一个子数组的最后一个元素与下一个子数组的第一个元素之差大于 20。

在上面的 arr 中,预期的 list 是 split_arr=[([5,6]),([28,29,32,33]),([87,88,95])]

回答1

就像我说的,如果 arr 看起来像:[5, 6, 16, 28, 29, 32, 33, 87, 88, 95],那么你的两个条件就有冲突,并且没有 num 16 的位置,所以下面的代码刚刚在条件1下,放弃条件2:

arr = [5, 6, 28, 29, 32, 33, 87, 88, 95]

results = []

idx = 0

sub_arr = list()

while idx <= len(arr) - 1:
    if not sub_arr:
        sub_arr.append(arr[idx])
    else:
        if arr[idx] - sub_arr[0] < 10:
            sub_arr.append(arr[idx])
        else:
            results.append(sub_arr)
            sub_arr = list()
            sub_arr.append(arr[idx])
    idx += 1

if sub_arr:
    results.append(sub_arr)

print(results)

输出:

[[5, 6], [28, 29, 32, 33], [87, 88, 95]]

回答2

它可以通过使用 NumPy 来实现,从条件 2 开始,然后是条件 1:

diff1 = np.diff(arr, prepend=arr[0]-21, append=arr[-1]+21)  # [21  1 22  1  3  1 54  1  7 21]
ind = np.where(diff1 > 20)[0]             # [0 2 6 9]
start = ind[:-1]                          # [0 2 6]
end = ind[1:] - 1                         # [1 5 8]
cond1 = (arr[end] - arr[start]) < 10      # [True True True]
result = np.split(arr, start[cond1][1:])  # [array([5, 6]), array([28, 29, 32, 33]), array([87, 88, 95])]

相似文章