如何将索引数组转换为Numpy中的遮罩数组?

39 投票
4 回答
35081 浏览
提问于 2025-04-19 17:02

有没有办法把一组索引转换成一组只有1和0的数组,给定一个范围?

比如说,输入[2,3],输出应该是[0, 0, 1, 1, 0],这个范围是5。

我想要自动化处理这样的事情:

>>> index_array = np.arange(200,300)
array([200, 201, ... , 299])

>>> mask_array = ???           # some function of index_array and 500
array([0, 0, 0, ..., 1, 1, 1, ... , 0, 0, 0])

>>> train(data[mask_array])    # trains with 200~299
>>> predict(data[~mask_array]) # predicts with 0~199, 300~499

4 个回答

1

根据你的要求,这里是答案。代码如下:

[x in index_array for x in range(500)]

这个代码会给你一个你想要的“遮罩”,不过它使用的是布尔值(也就是真和假),而不是0和1。

2

这里有个不错的技巧,可以用一行代码来实现这个功能,方法是使用 numpy.in1dnumpy.arange 这两个函数,像这样写(最后一行是关键部分):

>>> x = np.linspace(-2, 2, 10)
>>> y = x**2 - 1
>>> idxs = np.where(y<0)

>>> np.in1d(np.arange(len(x)), idxs)
array([False, False, False,  True,  True,  True,  True, False, False, False], dtype=bool)

不过,这种方法的缺点是速度大约比 Warren Weckesser 提供的方法慢 10 到 100 倍……但它只用一行代码,这可能正是你想要的。

14

对于一维的情况,可以试试:

n = (15,)
index_array = [2, 5, 7]
mask_array = numpy.zeros(n)
mask_array[index_array] = 1

如果是多维的情况,先把你的多维索引转换成一维的,然后再使用 ravel:

n = (15, 15)
index_array = [[1, 4, 6], [10, 11, 2]] # you may need to transpose your indices!
mask_array = numpy.zeros(n)
flat_index_array = np.ravel_multi_index(
    index_array,
    mask_array.shape)
numpy.ravel(mask_array)[flat_index_array] = 1
47

这里有一种方法:

In [1]: index_array = np.array([3, 4, 7, 9])

In [2]: n = 15

In [3]: mask_array = np.zeros(n, dtype=int)

In [4]: mask_array[index_array] = 1

In [5]: mask_array
Out[5]: array([0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0])

如果这个掩码总是一个范围的话,你可以省去 index_array,直接给一个切片赋值为 1

In [6]: mask_array = np.zeros(n, dtype=int)

In [7]: mask_array[5:10] = 1

In [8]: mask_array
Out[8]: array([0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0])

如果你想要一个布尔值的数组,而不是整数,可以在创建 mask_array 的时候改变它的 dtype

In [11]: mask_array = np.zeros(n, dtype=bool)

In [12]: mask_array
Out[12]: 
array([False, False, False, False, False, False, False, False, False,
       False, False, False, False, False, False], dtype=bool)

In [13]: mask_array[5:10] = True

In [14]: mask_array
Out[14]: 
array([False, False, False, False, False,  True,  True,  True,  True,
        True, False, False, False, False, False], dtype=bool)

撰写回答