如何在numpy中搜索满足条件的下标

Question

我在numpy数组'a'中有对应给定日期，月份和年份的列，我将这三个值与另一个名为'b'的数组'b'的列进行比较，该数组'b'也对应日期，月份和年份，以找到'a'等于'b'的索引到目前为止，我已经尝试过了:

a[:,3:6,1] == b[1,3:6]
array([[False,  True,  True],
   [ True,  True,  True],
   [False,  True,  True],
   ..., 
   [False, False, False],
   [False, False, False],
   [False, False, False]], dtype=bool)

可以正常工作，但我需要对应于[True,True,True]的行我也试过:

np.where(a[:,3:6,1] == b[1,3:6], a[:,3:6,1])
ValueError: either both or neither of x and y should be given

和

a[:,:,1].all(a[:,3:6,1] == b[1,3:6])
TypeError: only length-1 arrays can be converted to Python scalars

什么是快速和简单的方法来做到这一点?

Answer 1

您可以沿着最后一个轴使用np.all():

rows = np.where((a[:,3:6,1]==b[1,3:6]).all(axis=1))[0]

它将在rows中存储所有包含True值的行索引