在另一个2D数组中查找NumPy 2D数组的行的索引

2022-05-26 00:00:00 numpy numpy-ndarray

问题描述

我是《麻木》的新手。 我有2个2维阵列。我想在arr1中找到arr2的索引。请给我提建议。

    arr1 = [[1, 2, 3],
            [4, 5, 6],
            [7, 8, 9],
            [4, 5, 6],
            [1, 2, 3]]

    arr2 = [[1, 2, 3],
            [4, 5, 6],
            [7, 8, 9]]
    
    desired_output = [0, 1, 2, 1, 0]

解决方案

实现此目的的一种方法。

如果在中未找到arr1的任何行,则为简单起见,pos中的位置的值将为-1

这会大量使用NumPybroadcasting和indexing。请随时要求进一步澄清。

原始示例:

import numpy as np
arr1 = np.array([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9],
                 [4, 5, 6],
                 [1, 2, 3]])
arr2 = np.array([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9]])

inds = arr1 == arr2[:, None]
row_sums = inds.sum(axis = 2)
i, j = np.where(row_sums == 3) # Check which rows match in all 3 columns

pos = np.ones(arr1.shape[0], dtype = 'int64') * -1
pos[j] = i
pos
array([0, 1, 2, 1, 0])

示例2:

import numpy as np
arr1 = np.array([[1, 2, 4],
                 [4, 5, 6],
                 [7, 8, 9],
                 [4, 1, 6],
                 [1, 2, 3]])
arr2 = np.array([[1, 2, 3],
                 [4, 5, 6],
                 [7, 8, 9]])

inds = arr1 == arr2[:, None]
row_sums = inds.sum(axis = 2)
i, j = np.where(row_sums == 3)

pos = np.ones(arr1.shape[0], dtype = 'int64') * -1
pos[j] = i
pos
array([-1,  1,  2, -1,  0])

如果您有更多列数,只需将第i, j = np.where(row_sums == 3)行更改为i, j = np.where(row_sums == arr1.shape[1])

相关文章