用于多维数组的 Numpy Where 函数答案

【问题标题】：Numpy Where function for multidimensional arrays [duplicate]用于多维数组的 Numpy Where 函数
【发布时间】：2020-08-22 13:20:34
【问题描述】：

我有两个形状为 (N1, 2) 和 (N2, 2) 的数组：

import numpy as np
arr1 = np.array([[0, 1], [0, 2], [0, 3], [1, 2], [1, 3], [2, 3]])
arr2 = np.array([[0, 2], [0, 3], [1, 2], [1, 2]])

我知道 arr1 中的每个元素都是独一无二的。 arr2 中的元素不必是唯一的，arr2 中的每个元素在 arr1 中只出现一次。

我想要做的是，我想获得一个形状（N2,1）的索引数组，它只给出了 arr2 中元素在 arr1 中的位置。所以在这种情况下，我希望输出是

idxs = np.array([1, 2, 3, 3])

换句话说，我正在寻找类似 np.where 的东西，它适用于多维数组，但 np.where 没有轴关键字。

我可以选择类似的东西

idxs = []
for element in arr2: 
    for i, other_element in enumerate(arr1):
        if (element == other_element).all():
            idxs.append(i)
            break

但这很难看，我敢打赌有一个简洁的 numpy 解决方案。

【问题讨论】：

标签： python python-3.x numpy

【解决方案1】：

您可以使用broadcasting 来检查arr2 中每一行的成员资格，all 来检查所有元素在哪些行中匹配，argmax 来获取相应的索引：

((arr2[:,None] == arr1).all(2)).argmax(1)
# array([1, 2, 3, 3], dtype=int64)

【讨论】：