首页 > 解决方案 > 基于索引切割numpy数组

问题描述

我有一个一维numpy数组。此数组中两个后续值之间的差值要么是 1,要么大于 1。我想为每次出现差异大于一的情况将数组切割成部分。因此:

arr = numpy.array([77, 78, 79, 80, 90, 91, 92, 100, 101, 102, 103, 104])

应该成为

[array([77, 78, 79, 80]), array([90, 91, 92]), array([100, 101, 102, 103, 104])]

我有以下代码可以解决问题,但我觉得我在这里很复杂。必须有更好/更蟒蛇的方式。有更优雅的方法的人吗?

import numpy

def split(arr, cut_idxs):

  empty_arr = []
  for idx in range(-1, cut_idxs.shape[0]):  
    if idx == -1:
      l, r = 0, cut_idxs[0]
    elif (idx != -1) and (idx != cut_idxs.shape[0] - 1):
      l, r = cut_idxs[idx] + 1, cut_idxs[idx + 1]
    elif idx == cut_idxs.shape[0] - 1:
      l, r = cut_idxs[-1] + 1, arr.shape[0]

    empty_arr.append(arr[l:r + 1]) 

  return empty_arr 


arr = numpy.array([77, 78, 79, 80, 90, 91, 92, 100, 101, 102, 103, 104])
cuts = numpy.where(numpy.ediff1d(arr) > 2)[0]

print split(arr, cuts)

标签: pythonarraysnumpyoptimization

解决方案


一种 Pythonic 方式是 -

np.split(arr, np.flatnonzero(np.diff(arr)>1)+1)

样品运行 -

In [10]: arr
Out[10]: array([ 77,  78,  79,  80,  90,  91,  92, 100, 101, 102, 103, 104])

In [11]: np.split(arr, np.flatnonzero(np.diff(arr)>1)+1)
Out[11]: 
[array([77, 78, 79, 80]),
 array([90, 91, 92]),
 array([100, 101, 102, 103, 104])]

另一个与slicing-

In [16]: cut_idx = np.r_[0,np.flatnonzero(np.diff(arr)>1)+1,len(arr)]
             # Or np.flatnonzero(np.r_[True, np.diff(arr)>1, True])

In [17]: [arr[i:j] for i,j in zip(cut_idx[:-1],cut_idx[1:])]
Out[17]: 
[array([77, 78, 79, 80]),
 array([90, 91, 92]),
 array([100, 101, 102, 103, 104])]

推荐阅读