Python 对具有滞后的列求和的最快方法

给定一个方阵，我想用它的行号移动每一行，并对列求和。例如：

array([[0, 1, 2],        array([[0, 1, 2],
       [3, 4, 5],    ->            [3, 4, 5],      ->   array([0, 1+3, 2+4+6, 5+7, 8]) = array([0, 4, 12, 12, 8])
       [6, 7, 8]])                    [6, 7, 8]])

我有4种解决方案-、

慢

和

最慢

，它们做的事情完全相同，并且按速度排列：

def fast(A):
    n = A.shape[0]
    retval = np.zeros(2*n-1)
    for i in range(n):
        retval[i:(i+n)] += A[i, :]
    return retval

令人惊讶的是，非矢量化的解决方案是最快的。以下是我的基准：

A = np.random.randn(1000,1000)

%timeit fast(A)
# 1.85 ms ± 20 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit slow(A)
# 3.28 ms ± 9.55 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit slower(A)
# 4.07 ms ± 18.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit slowest(A)
# 58.4 ms ± 993 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

---

是否存在更快的解决方案？如果不是的话，有人能解释一下为什么事实上

fast

是最快的吗

编辑

对于

慢速

，稍微加速：

def slow(A):
    n = A.shape[0]
    indices = np.arange(2*n-1)
    indices = np.lib.stride_tricks.as_strided(indices, A.shape, (8,8))
    return np.bincount(indices.ravel(), A.ravel())

以与Pierre相同的方式绘制运行时（以2**15为上限-出于某种原因

慢速
无法处理此大小）

对于
100x100
阵列，
slow
比任何解决方案（不使用
numba
）都要稍微快一些
sum\u反对角线仍然是
1000x1000
阵列的最佳选择。
最简单的加速方法（在我的电脑上~2/3次）是使用
fast
方法和
numba
软件包：

import numba
@numba.jit(nopython=True)
def fastNumba(A):
    n = A.shape[0]
    retval = np.zeros(2*n-1)
    for i in range(n):
        retval[i:(i+n)] += A[i, :]
    return retval

但是只有当此函数运行多次时，使用
numba
才有意义。函数的第一次求值需要更多的时间（因为编译）。
这里有一种方法，有时比您的“
fast（）
”版本快，但对于
nxn
数组，它只在
n
的有限范围内（大约在30到1000之间）。循环（
fast（）
）在大型数组上很难击败，即使使用
numba
，实际上也会渐近收敛到简单的
a.sum（axis=0）
，这表明这与大型数组的效率差不多（感谢您的循环！）

我将调用的方法是
sum\u antidigonals（）
，它使用
np.add.reduce（）
对
a
的条带化版本和从一个相对较小的1D数组生成的合成掩码进行条带化，以创建2D数组的幻觉（无需消耗更多内存）

此外，它不仅限于方形数组（但是
fast（）
也可以很容易地适应这种泛化，请参阅本文底部的
fast\u g（）
）

def sum_反对角线（a）：
断言a.flags.c_连续
r、 c=a.形状
s0，s1=a.跨步
z=np.lib.stride\u.as\u(
a、 形状=（r，c+r-1），步幅=（s0-s1，s1），可写=假）
#面具
kern=np.r_uu[np.repeat（False，r-1），np.repeat（True，c），np.repeat（False，r-1）]
mask=np.fliplr（np.lib.stride\u tricks.as\u stride(
kern，shape=（r，c+r-1），跨步=（1，1），writeable=False）
返回np.add.reduce（z，其中=掩码）

请注意，它不限于方形阵列：

求和反对角线（np.arange（15）.重塑（5,3））
数组（[0,4,12,21,30,24,14]）

解释

为了理解它是如何工作的，让我们用一个示例来检查这些条带化阵列

给定一个初始数组
a
，即
（3，2）
：

a=np.arange（6）.重塑（3,2）
>>>a
数组（[[0，1]，
[2, 3],
[4, 5]])
#在计算函数中的z之后
>>>z
数组（[[0,1,2,3]，
[1, 2, 3, 4],
[2, 3, 4, 5]])

您可以看到，它几乎就是我们想要的总和（axis=0）
，除了上面和下面的三角形是不需要的。我们真正想要总结的是：

array([[0, 1, -, -],
       [-, 2, 3, -],
       [-, -, 4, 5]])

输入掩码，我们可以从1D内核开始构建它：

kern=np.r\n[np.repeat（False，r-1），np.repeat（True，c），np.repeat（False，r-1）]
>>>克恩
数组（[False，False，True，True，False，False]）

我们使用了一个有趣的片段：
（1，1）
，这意味着我们重复同一行，但每次滑动一个元素：

>>np.lib.stride\u tricks.as\u(
…kern，shape=（r，c+r-1），strips=（1，1），writeable=False）
数组（[[False，False，True，True]，
[假，真，真，假]，
[对，对，错，错]]

然后我们只需向左/向右翻转它，并将其用作
np.add.reduce（）
的
where
参数

速度

b=np.random.normal（大小=（10001000））
#检查与OP的fast（）函数的等价性：
>>>np.allclose（快速（b），求和反对角线（b））
真的
%时间和反对角线（b）
#每个回路1.83 ms±840 ns（7次运行的平均值±标准偏差，每个1000个回路）
%快速计时（b）
#每个回路2.07 ms±15.2µs（7次运行的平均值±标准偏差，每个100个回路）

在这种情况下，速度稍微快一点，但只有10%左右

在300x300阵列上，
sum\u antidigonals（）
比
fast（）
快2.27倍

然而

即使将
z
和
mask
组合起来非常快（在上面的1000x1000示例中，
np.add.reduce（）
之前的整个设置只需要46µs），总和本身也是
O[r（r+c）]
，即使只需要
O[r c]
实际添加（其中
mask==True
）。因此，对于方形阵列，要执行的操作大约要多2倍

在10K x 10K阵列上，我们发现：


fast
需要95毫秒，而
sum_反对角线
需要208毫秒

通过大小范围进行比较

我们将使用可爱的软件包在
n
范围内比较多种方法的速度：

perfplot.show(
setup=lambda n:np.random.normal（大小=（n，n）），
内核=[just_sum_0，fast，fast_g，nb_fast_i，nb_fast_ij，sum_反对角线]，
n_范围=[2**k表示范围（3,16）中的k]，
相等性检查=无，因为只有和0
xlabel='n'，
相对_to=1，
)


观察结果


如您所见，
sum\u antidagonals（）
相对于
fast（）
的速度优势被限制在
n
的范围内，大约在30到1000之间
它永远比不上
def slow(A):
    n = A.shape[0]
    indices = np.arange(2*n-1)
    indices = np.lib.stride_tricks.as_strided(indices, A.shape, (8,8))
    return np.bincount(indices.ravel(), A.ravel())

import numba
@numba.jit(nopython=True)
def fastNumba(A):
    n = A.shape[0]
    retval = np.zeros(2*n-1)
    for i in range(n):
        retval[i:(i+n)] += A[i, :]
    return retval

array([[0, 1, -, -],
       [-, 2, 3, -],
       [-, -, 4, 5]])