如何快速获取numpy数组中的非零值索引?

时间:2016-06-27 09:02:07

标签: python algorithm numpy data-structures

现在我正在编写一个函数来获取非零值的索引,并遵循以下规则:

  1. 预期结果是一个列表。每个元素表示非零值的连续切片的索引。因此,对于[0,0,0,1,1,1,0,1,1,0]的列表,它应该获得列表[[3,4,5], [7,8]]
  2. 列表中不同值的索引应位于单独的列表中,即对于[0,0,1,1,1,2,2,1,1,0]列表,预期结果为[[2,3,4],[5,6],[7,8]]
  3. 你知道吗?先感谢您!

1 个答案:

答案 0 :(得分:3)

使用arr作为输入数组并将数组列表作为输出,你可以这样做 -

# Store non-zero element indices
idx = np.where(arr)[0]

# Get indices where the shifts occur, i.e. positions where groups of identical 
# elements are separated. For this we perform differnetiation and look for 
# non-zero values and then get those positions. Finally, add 1 to compensate 
# for differentiation that would have decreased those shift indices by 1.
shift_idx = np.where(np.diff(arr[idx])!=0)[0]+1

# Split the non-zero indices at those shifts for final output
out = np.split(idx,shift_idx)

示例输入,输出 -

In [35]: arr
Out[35]: array([0, 0, 1, 1, 1, 2, 2, 1, 1, 0, 2, 2, 4, 3, 3, 3, 0])

In [36]: out
Out[36]: 
[array([2, 3, 4]),
 array([5, 6]),
 array([7, 8]),
 array([10, 11]),
 array([12]),
 array([13, 14, 15])]