如何在 numpy 中搜索数组的每一行中的索引值创建一个布尔数组

Question

给定一个大小为 MxN 的数组和一个大小为 Mx1 的数组，我想用 MxN 计算一个布尔数组。

import numpy as np

M = 2

N = 3

a = np.random.rand(M, N) # The values doesn't matter

b = np.random.choice(a=N, size=(M, 1), replace=True)

# b =

# array([[2],

#        [1]])

# I found this way to compute the boolean array but I wonder if there's a fancier, elegant way

index_array = np.array([np.array(range(N)), ]*M)

# Create an index array

# index_array = 

# array([[0, 1, 2],

#        [0, 1, 2]])

#

boolean_array = index_array == b

# boolean_array =

# array([[False, False,  True],

#        [False,  True, False]])

#

所以我想知道是否有一种更奇特的 pythonic 方式来做到这一点

Answer 1

您可以通过直接广播与单个 1d 范围的比较来简化：M = 2N = 3a = np.random.rand(M, N) b = np.random.choice(a=N, size=(M, 1), replace=True)print(b)array([[1],       [2]])b == np.arange(N)array([[False,  True, False],       [False, False,  True]])通常，广播在这些情况下很方便，因为它使我们不必创建形状兼容的数组来执行与其他数组的操作。对于生成的数组，我可能会改为使用以下内容：np.broadcast_to(np.arange(N), (M,N))array([[0, 1, 2],       [0, 1, 2]])尽管如前所述，NumPy 使这里的生活更轻松，因此我们不必为此担心。