将 numpy 布尔选择器限制为前几个真值的有效方法答案

【问题标题】：Efficient way to restrict a numpy boolean selector to the first few true values将 numpy 布尔选择器限制为前几个真值的有效方法
【发布时间】：2019-11-21 11:47:02
【问题描述】：

我有一个 numpy boolean selector 数组，我可以将其应用于数组 a。 （实际上在问题域中不是随机的，这只是为了方便示例）。但我实际上只想使用 selector 的前 n 个 True 条目进行选择（示例中最多为 n=3）。那么给定selector加上参数n，如何使用numpy操作生成select_first_few，从而避免迭代循环？

>>> import numpy as np
>>> selector = np.random.random(10) > 0.5
>>> a = np.arange(10)
>>> selector
array([ True, False,  True,  True,  True, False,  True, False,  True,
       False])
>>> chosen, others = a[selector], a[~selector]
>>> chosen
array([0, 2, 3, 4, 6, 8])
>>> others
array([1, 5, 7, 9])
>>> select_first_few = np.array([ True, False,  True,  True,  False, False,  False, False,  False,
...        False])
>>> chosen_few, tough_luck = a[select_first_few], a[~select_first_few]
>>> chosen_few
array([0, 2, 3])
>>> tough_luck
array([1, 4, 5, 6, 7, 8, 9])

【问题讨论】：

切片结果？
我可以看到它如何生成 chosen_few 但不能生成 tough_luck

标签： python numpy boolean

【解决方案1】：

方法#1

一种方法是使用 cumsum 和 argmax 获取范围，然后切片以设置 False -

In [40]: n = 3

In [41]: selector
Out[41]: 
array([ True, False,  True,  True,  True, False,  True, False,  True,
       False])

In [42]: selector[(selector.cumsum()>n).argmax():] = 0

In [43]: selector # your select_first_few mask
Out[43]: 
array([ True, False,  True,  True, False, False, False, False, False,
       False])

然后，使用这个新的selector 从输入数组中选择和取消选择元素。

方法 #2

另一种方法是 mask-the-mask -

n = 3
C = np.count_nonzero(selector)
newmask = np.zeros(C, dtype=bool)
newmask[:n] = 1
selector[selector] = newmask

示例运行 -

In [62]: selector
Out[62]: 
array([ True, False,  True,  True,  True, False,  True, False,  True,
       False])

In [63]: n = 3
    ...: C = np.count_nonzero(selector)
    ...: newmask = np.zeros(C, dtype=bool)
    ...: newmask[:n] = 1
    ...: selector[selector] = newmask

In [64]: selector
Out[64]: 
array([ True, False,  True,  True, False, False, False, False, False,
       False])

或者通过布尔值的即时连接使其更短 -

n = 3
C = np.count_nonzero(selector)
selector[selector] = np.r_[np.ones(n,dtype=bool),np.zeros(C-n,dtype=bool)]

方法#3

最简单的一个 -

selector &= selector.cumsum()<=n

【讨论】：

方法#1 selector[(selector.cumsum()>n).argmax():] = 0 几乎正是我想要的。一个警告：如果 n 超过 selector 中可用 True 值的数量，例如 n=8，则此表达式将其更改为所有 False，而在这种情况下我可能希望 selector 保持不变。
@DavidWallace 你可以选择刚刚添加的 - Approach #3 在这种情况下。

【解决方案2】：

获取列表中所有选择的索引并对该列表进行切片。然后使用列表推导来检索那些选定索引处的数据。

import numpy as np
selector = np.random.random(10) > 0.5
data = np.arange(10)

choosen_indices = np.where(selector)

#select first 3 choosen
choosen_few_indices = choosen_indices[:3]
choosen_few = [data[i] for i in choosen_few_indices]

# if you are also interested in the not choosen data
not_choosen_indices = list(set(range(len(data))) - set(choosen_indices))
# proceed ...

【讨论】：