在可变距离上高效的 NumPy 行旋转答案

【问题标题】：Efficient NumPy rows rotation over variable distances在可变距离上高效的 NumPy 行旋转
【发布时间】：2016-11-25 02:49:12
【问题描述】：

给定一个 2D M x N NumPy 数组和一个旋转距离列表，我想在列表中的距离上旋转所有 M 行。这是我目前拥有的：

import numpy as np

M = 6
N = 8
dists = [2,0,2,1,4,2] # for example
matrix = np.random.randint(0,2,(M,N))

for i in range(M):
    matrix[i] = np.roll(matrix[i], -dists[i])

最后两行实际上是一个内部循环的一部分，它被执行了数十万次，它正在成为我通过 cProfile 衡量的性能的瓶颈。例如，是否有可能避免 for 循环并更有效地做到这一点？

【问题讨论】：

标签： python python-3.x numpy optimization

【解决方案1】：

在将dists 与range(0...N) 数组相加后，我们可以通过取模操作模拟滚动行为，从而为每一行提供在同一行中挑选和洗牌的元素的列索引。在broadcasting 的帮助下，我们可以将这个过程矢量化到所有行。因此，我们会有这样的实现 -

M,N = matrix.shape # Store matrix shape

# Get column indices for all elems for a rolled version with modulus operation
col_idx = np.mod(np.arange(N) + dists[:,None],N)

# Index into matrix with ranged row indices and col indices to get final o/p
out = matrix[np.arange(M)[:,None],col_idx]

【讨论】：

虽然当我使用 timeit 模块测试这些行时，这几乎快了 65%，但当我将它放入总代码中时，它实际上使它慢了 20% 慢（多次运行的平均值，不是那么高的标准偏差）。任何猜测为什么会这样？
@GekkeBoer 在您的实际情况下，矩阵的形状是什么？
我有更多案例，但我刚才在用 4x4 进行基准测试。但是，通过您的方法，我可以在一些循环之外计算 col_idx，然后它会快很多！谢谢！我会稍等片刻接受这个作为正确答案，以防有人仍然能够改进它。
@GekkeBoer 是的，如果您正在使用相同形状的数组和dists，这个col_idx 应该可以重复使用。