在一个二维的numpy数组中，针对每一行获取在另一个二维数组中相等行的索引。

Question

我有两个巨大的2d numpy整数数组X和U，其中假定U只有唯一的行。对于X中的每一行，我想要获取与U中匹配的行的相应行索引（如果有，则为-1）。例如，如果将以下数组作为输入传递：

U = array([[1, 4],
       [2, 5],
       [3, 6]])

X = array([[1, 4],
       [3, 6],
       [7, 8],
       [1, 4]])

输出应该是：

array([0,2,-1,0])

有没有一种有效的方法用Numpy做这个（或类似的）？@Divakar：你的方法对我无效。

print(type(rows), rows.dtype, rows.shape)
print(rows[:10])
print(search2D_indices(rows[:10], rows[:10]))

<class 'numpy.ndarray'> int32 (47398019, 5)
[[65536     1     1     1    17]
 [65536     1     1     1   153]
 [65536     1     1     2   137]
 [65536     1     1     3   153]
 [65536     1     1     9   124]
 [65536     1     1    13   377]
 [65536     1     1    13   134]
 [65536     1     1    13   137]
 [65536     1     1    13   153]
 [65536     1     1    13   439]]
[ 0  1  2  3  4 -1 -1 -1 -1  9]

- user2224350

Answer 1

方法一

受this solution启发，用于查找numpy数组中多个值的行索引的向量化解决方案是使用searchsorted -

def search2D_indices(X, searched_values, fillval=-1):
    dims = np.maximum(X.max(0), searched_values.max(0))+1
    X1D = np.ravel_multi_index(X.T,dims)
    searched_valuesID = np.ravel_multi_index(searched_values.T,dims)
    sidx = X1D.argsort()
    idx = np.searchsorted(X1D,searched_valuesID,sorter=sidx)
    idx[idx==len(sidx)] = 0    
    idx_out = sidx[idx]
    return np.where(X1D[idx_out] == searched_valuesID, idx_out, fillval)

样例运行 -

In [121]: U
Out[121]: 
array([[1, 4],
       [2, 5],
       [3, 6]])

In [122]: X
Out[122]: 
array([[1, 4],
       [3, 6],
       [7, 8],
       [1, 4]])

In [123]: search2D_indices(U, X, fillval=-1)
Out[123]: array([ 0,  2, -1,  0])

---

方法二

针对负整数的情况，我们需要相应地偏移dims和转换为1D，如下所示 -

def search2D_indices_v2(X, searched_values, fillval=-1):
    X_lim = X.max()-X.min(0)
    searched_values_lim = searched_values.max()-searched_values.min(0)

    dims = np.maximum(X_lim, searched_values_lim)+1
    s = dims.cumprod()

    X1D = X.dot(s)
    searched_valuesID = searched_values.dot(s)
    sidx = X1D.argsort()
    idx = np.searchsorted(X1D,searched_valuesID,sorter=sidx)
    idx[idx==len(sidx)] = 0    
    idx_out = sidx[idx]

    return np.where(X1D[idx_out] == searched_valuesID, idx_out, fillval)

示例运行 -

In [142]: U
Out[142]: 
array([[-1, -4],
       [ 2,  5],
       [ 3,  6]])

In [143]: X
Out[143]: 
array([[-1, -4],
       [ 3,  6],
       [ 7,  8],
       [-1, -4]])

In [144]: search2D_indices_v2(U, X, fillval=-1)
Out[144]: array([ 0,  2, -1,  0])

方法三

另一种基于 视图 的方法 -

# https://dev59.com/tqPia4cB1Zd3GeqP3cEw#45313353/ @Divakar
def view1D(a, b): # a, b are arrays
    a = np.ascontiguousarray(a)
    b = np.ascontiguousarray(b)
    void_dt = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
    return a.view(void_dt).ravel(),  b.view(void_dt).ravel()

def search2D_indices_views(X, searched_values, fillval=-1):
    X1D,searched_valuesID = view1D(X, searched_values)
    sidx = X1D.argsort()
    idx = np.searchsorted(X1D,searched_valuesID,sorter=sidx)
    idx[idx==len(sidx)] = 0    
    idx_out = sidx[idx]
    return np.where(X1D[idx_out] == searched_valuesID, idx_out, fillval)

- Divakar

---

这种方法对我不起作用，请看一下我编辑过的问题。 - user2224350

@user2224350 是的，确实是个bug。刚刚做的修改应该已经修复了它。 - Divakar

嗯，我仍然得到相同的错误。不过第三种方法有效。这真的很快！非常感谢！ - user2224350

可能是溢出的情况。尝试将输入转换为 np.int64 数据类型。 - Divakar

Answer 2

这是一种基于字典的方法：

import numpy as np

U = np.array([[1, 4],
              [2, 5],
              [3, 6]])

X = np.array([[1, 4],
              [3, 6],
              [7, 8],
              [1, 1]])

d = {v: k for k, v in enumerate(map(tuple, U))}

res = np.array([d.get(tuple(a), -1) for a in X])

# [ 0  2 -1 -1]

- jpp

---

我猜这会非常慢，因为有for循环。X和U大约有几千万到一亿行。 - user2224350

Answer 3

你可以使用广播（broadcasting）以向量化的方式确定向量中各项的公平性。然后，你可以简单地使用all函数在适当的轴上获取所需的真值，对应于预期的索引。最后，使用np.where获取公平性发生的索引，并将其重新分配给先前创建的填充有-1的数组。

In [47]: result = np.full(X.shape[0], -1)

In [48]: x, y = np.where((X[:,None] == U).all(-1))

In [49]: result[x] = y

In [50]: result
Out[50]: array([ 0,  2, -1,  0])

请注意，正如文档中所提到的那样，在广播方面，您应该注意以下几点：

尽管这在代码行方面非常高效，但在计算效率方面可能并不高效。问题在于算法的中间步骤中计算的三维差分数组。对于小数据集，创建和操作数组很可能非常快。然而，大数据集将生成一个大的中间数组，这在计算效率上是低效的。

- Mazdak

---

为什么要踩票？这个答案有什么问题我可以改进吗？ - Mazdak

@Kasramvd 不是我，我同意jpp的观点！ - Divakar

@jpp 这个解决方案完全向量化，但肯定不是最好的/实际上没有最好的lol。但是尽管如此，这违反了SO的投票规则。只有在答案错误或与给定问题无关时才应该投反对票。 - Mazdak

@Kasramvd 不是我点的踩，但就我看来，我真的不认为这“违反了SO的踩规则”。你可以随意踩。而当OP的问题的第一行是“我有两个巨大的NumPy数组”时，你的解决方案看起来并不可行。 - miradulo

@miradulo，它确实有用，但这不是重点。你说得对，“你可以对任何你想要的东西进行投票”，这是一个显而易见的事实，你确实可以，但这并不意味着你应该这样做。另一点是，这个解决方案比其他基于Numpy的方法更直接和易懂（尽管没有冒犯的意思），并且正确回答了问题，但不是以惊人的速度。这段代码有许多优点，其中之一可能是以向量化的方式解决问题，还有广播的优缺点。 - Mazdak