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

【问题标题】：Most efficient way to rearrange 2D numpy array from another 2D index array从另一个二维索引数组重新排列二维 numpy 数组的最有效方法
【发布时间】：2020-10-25 07:24:39
【问题描述】：

简述

在 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)))成功地将A使用B作为索引数组重新排列：

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)))

【问题讨论】：

np.take_along_axis(A,B,1)?
A[ np.arange(5)[:,None],B] 也应该可以工作，但take_along 更容易（如果你记得它存在的话:)）。
@PaulPanzer 我做了一些测试，take_along_axis 函数实际上比 FOR 循环慢。谜团……
糟糕！你的阵列很小吗？ @hpaulj 的建议呢？
@PaulPanzer 哦，之前评论的不是我（OP）。我的数组可能相当大，比 20000 x 10000 大得多。我正在使用@bousof 的建议，看起来循环对于大型 nCol 来说是最有吸引力的。随着 nCols 的减少，take_along_axis 和 @hpaulj 的速度更快

标签： python arrays performance numpy

【解决方案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

【讨论】：