python - 从另一个二维索引数组重新排列二维 numpy 数组的最有效方法

Question

简述

在 Python 3.6 中并使用 Numpy，根据不同的、形状相似的索引二维数组中存在的索引重新排列二维数组元素的最有效方法是什么？

详细

假设我有以下两个 9 x 5 数组，分别称为 A 和 B:

import numpy as np
A = np.array([[0.32, 0.35, 0.88, 0.63, 1.  ],
              [0.23, 0.69, 0.98, 0.22, 0.96],
              [0.7 , 0.51, 0.09, 0.58, 0.19],
              [0.98, 0.42, 0.62, 0.94, 0.46],
              [0.48, 0.59, 0.17, 0.23, 0.98]])

B = np.array([[4, 0, 3, 2, 1],
              [3, 2, 4, 1, 0],
              [4, 3, 0, 2, 1],
              [4, 2, 0, 3, 1],
              [0, 3, 1, 2, 4]])

我可以通过 np.array(list(map(lambda i, j: j[i], B, A))) 使用 B 作为索引数组成功地重新排列 A:

array([[1.  , 0.32, 0.63, 0.88, 0.35],
       [0.22, 0.98, 0.96, 0.69, 0.23],
       [0.19, 0.58, 0.7 , 0.09, 0.51],
       [0.46, 0.62, 0.98, 0.94, 0.42],
       [0.48, 0.23, 0.59, 0.17, 0.98]])

但是，当 A 和 B 的维度增加时，这种解决方案变得非常低效。如果我没记错的话，那是因为:

• 对 A 的所有行使用 lambda 循环，而不是依赖 Numpy 向量化
• 映射很慢
• 将列表转换为数组占用了宝贵的时间。

因为在我的实际用例中，这些数组可能会变得很大，而且我必须在一个长循环中对其中的许多数组重新排序，所以我当前的很多性能瓶颈(使用分析器测量)都来自上面的那一行代码.

我的问题:实现上述目标的最有效、更 Numpy 智能的方法是什么？

测试一般数组和时间过程的玩具代码可以是:

import numpy as np
nRows = 20000
nCols = 10000
A = np.round(np.random.uniform(0, 1, (nRows, nCols)), 2)
B = np.full((nRows, nCols), range(nCols))
for r in range(nRows):
    np.random.shuffle(B[r])
%time X = np.array(list(map(lambda i, j: j[i], B, A)))

Answer 1

与其他三种可能性的比较:

import numpy as np
import time

# Input
nRows = 20000
nCols = 10000
A = np.round(np.random.uniform(0, 1, (nRows, nCols)), 2)
B = np.full((nRows, nCols), range(nCols))
for r in range(nRows):
  np.random.shuffle(B[r])

# Original
t_start = time.time()
X = np.array(list(map(lambda i, j: j[i], B, A)))
print('Timer 1:', time.time()-t_start, 's')

# FOR loop
t_start = time.time()
X = np.zeros((nRows, nCols))
for i in range(nRows):
  X[i] = A[i][B[i]]
print('Timer 2:', time.time()-t_start, 's')

# take_along_axis
t_start = time.time()
X = np.take_along_axis(A,B,1)
print('Timer 3:', time.time()-t_start, 's')

# Indexing
t_start = time.time()
X = A[ np.arange(nRows)[:,None],B]
print('Timer 4:', time.time()-t_start, 's')

输出:

% python3 script.py
Timer 1: 2.191567897796631 s
Timer 2: 1.3516249656677246 s
Timer 3: 1.675267219543457 s
Timer 4: 1.646852970123291 s

对于较少的列数 (nRows,nCols)=(200000,10) 然而，结果完全不同:

% python3 script.py
Timer 1: 0.2729799747467041 s
Timer 2: 0.22678399085998535 s
Timer 3: 0.016162633895874023 s
Timer 4: 0.014748811721801758 s