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 API 执行此操作的最简单方法有什么建议吗？学习 python 和 numpy，将不胜感激。谢谢!

Answer 1

如果您这样设置问题:

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

等等