从掩码中获取切片索引 Python

Question

我有一个 (N,) 浮点数数组 (arr)，但我只关心 >= 给定 threshold 的条目。我可以获得这样的面具：

mask = (arr >= threshold)

现在我想要一个 (N,2) 对应切片索引的数组。

例如，如果 arr = [0, 0, 1, 1, 1, 0, 1, 1, 0, 1] 和 threshold = 1，则 mask = [False, False, True, True, True, False, True, True, False, True]，我想要索引 [ [2, 5], [6, 8], [9, 10] ]（我可以将其用作 arr[2:5], arr[6:8], arr[9:10]获取 arr >= threshold).

的段

目前，在将相应的切片索引附加到列表之前，我有一个丑陋的 for 循环解决方案，它遵循 True 的每一段。有没有更简洁易读的方法来实现这个结果？

Answer 1

您可以使用带 key 参数的 itertools groupby 以及 enumerate 来获取分组。如果组值都是 True 你可以取 first 和 last+1 值。

from itertools import groupby
import numpy as np
arr = np.array([0, 0, 1, 1, 1, 0, 1, 1, 0, 1])
threshold  = 1


idx = []
for group,data in groupby(enumerate((arr >= threshold)), key=lambda x:x[1]):
    d = list(data)
    if all(x[1]==True for x in d):
        idx.append([d[0][0], d[-1][0]+1])

输出

[[2, 5], [6, 8], [9, 10]]

Answer 2

您可以使用 np.flatnonzero and np.diff:

的组合

indexes = np.flatnonzero(np.diff(np.append(arr >= threshold, 0))) + 1
indexes = list(zip(indexes[0::2], indexes[1::2]))

输出：

>>> indexes
[(2, 5), (6, 8), (9, 10)]

Answer 3

您可以使用掩码通过将掩码布尔值与其后继者进行比较来计算开始索引和结束索引的列表。然后连接开始和结束以形成范围（全部使用 numpy 方法矢量化）：

import numpy as np

arr       = np.array([0, 0, 1, 1, 1, 0, 1, 1, 0, 1])
threshold = 1
mask      = arr >= threshold

starts    = np.argwhere(np.insert(mask[:-1],0,False)<mask)[:,0]
ends      = np.argwhere(np.append(mask[1:],False)<mask)[:,0]+1
indexes   = np.stack((starts,ends)).T

print(starts)  # [2 6 9]
print(ends)    # [5 8 10]
print(indexes)
[[ 2  5]
 [ 6  8]
 [ 9 10]]

如果您想要 Python 元组列表中的结果：

indexes = list(zip(starts,ends))  # [(2, 5), (6, 8), (9, 10)]

如果您不需要（或不想）使用 numpy，您可以使用 itertools 中的 groupby 直接从 arr 获取范围：

from itertools import groupby

indexes = [ (t[1],t[-1]+1) for t,t[1:] in 
            groupby(range(len(arr)),lambda i:[arr[i]>=threshold]) if t[0]]

print(indexes)
[(2, 5), (6, 8), (9, 10)]

从掩码中获取切片索引 Python

Get slice indices from mask Python

python

numpy

slice