Python - 切片数组直到满足某条件

Question

我需要从一个数组的某个索引开始切割，直到满足某个条件为止。

>>> a = numpy.zeros((10), dtype='|S1')
>>> a[2] = 'A'
>>> a[4] = 'X'
>>> a[8] = 'B'
>>> a
array(['', '', 'A', '', 'X', '', '', '', 'B', ''], dtype='|S1')

比如，对于上面的数组，我想从某个索引开始，切割到第一个非零值为止，向前和向后都要考虑。例如，对于索引值2、4、8，结果会是：

['', '', A, '']      # 2
['', X, '', '', '']  # 4
['', '', '', B, '']  # 8

有没有简单的方法可以用numpy这个库来实现这个？我正在学习python和numpy，任何帮助都非常感谢！

Answer 1

注意，这个可以通过纯Python的itertools和functools模块来简单实现。

import functools, itertools
arr = ['', '', 'A', '', 'X', '', '', '', 'B', '']

f = functools.partial(itertools.takewhile, lambda x: not x)
def g(a, i):
    return itertools.chain(f(reversed(a[:i])), [a[i]], f(a[i+1:]))

我们把f定义为一个子迭代器，它会一直查找，直到找到一个为真的元素。而g则是将这个方法应用在列表中索引前面的部分和索引后面的部分的组合。

这样会返回生成器，我们可以把它们转换成包含我们结果的列表。

>>> list(g(arr, 2))
['', '', 'A', '']
>>> list(g(arr, 4))
['', 'X', '', '', '']
>>> list(g(arr, 8))
['', '', '', 'B', '']

Answer 2

如果你把问题设置成这样：

import numpy
a = numpy.zeros((10), dtype=str)
a[2] = 'A'
a[4] = 'X'
a[8] = 'B'

你可以很简单地找到非空字符串的索引，方法如下：

i = numpy.where(a!='')[0]  # array([2, 4, 8])

另外，numpy.argwhere(..) 也能很好地解决这个问题。

然后你可以用这个数组来进行切片操作：

out2 = a[:i[1]]        # 2   ['' '' 'A' '']
out4 = a[i[0]+1:i[2]]  # 4   ['' 'X' '' '' '']

等等。

Answer 3

这是关于掩码数组的工作，numpy.ma有很多函数可以处理子集。

a = np.zeros((10), dtype=str)
a[2] = 'A'
a[4] = 'X'
a[8] = 'B'

让我们把非空的元素掩盖掉：

am=np.ma.masked_where(a!='', a)

np.ma.notmasked_contiguous会高效地遍历数组，找到所有未被掩盖的连续元素片段：

slices = np.ma.notmasked_contiguous(am)
[slice(0, 1, None), slice(3, 3, None), slice(5, 7, None), slice(9, 9, None)]

比如说，数组在第5个和第7个元素之间是连续空的。 现在你只需要连接你感兴趣的片段，首先获取每个片段的起始索引：

slices_start = np.array([s.start for s in slices])

然后找到你想要的索引的位置：

slices_start.searchsorted(4) #4
Out: 2

所以你想要第1个和第2个片段： a[slices[1].start:slices[2].stop+1] array(['', 'X', '', '', ''], dtype='|S1')

或者我们试试8：

i = slices_start.searchsorted(8)
a[slices[i-1].start:slices[i].stop+1]
Out: array(['', '', '', 'B', ''], 
  dtype='|S1')

你可能应该在ipython里多玩玩这个，以更好地理解它。