时间戳

首页 > 解决方案 > numpy：以移位的方式将数组相加

python - numpy：以移位的方式将数组相加

问题描述

假设我有一个形状为k*k的n 个数组的输入。现在我想将它们加在一起，但在每次操作后将k*k数组移动一些正整数（步长） 。

很难真正解释它，所以这里有一张小图。

在此处输入图像描述

在 k*k 数组相交的每个部分，它们只是被相加。

我目前的实现方式如下：

import numpy as np

# some input with shape (n ,k ,k)
x = np.arange(16).reshape(4 ,2 ,2)
# im intializing the output to zeros in the beginning and
# add the k*k windows to it
out = np.zeros((2 ,5))

# loops over the (n ,k ,k) input and adds the values
# to the output and shifts by 1 to the right
for i in range(4):
   # i could also be multipied by any integer to move 
   # the window by more then one 
   out[: ,i+0:i+2] += x[i]

所以我的问题是是否有可能摆脱循环，或者您是否可以使用一些内置的 numpy 函数来加速这个或更通用的方法，以便在输入大小增加时使其更快。

感谢您的阅读，祝您有愉快的一天:)

标签： pythonarraysnumpy

解决方案

我能想到几种方法。最暴力的是创建一个(n, k, k + steps * (n - 1))数组，然后对第一个轴求和：

steps = ...
n, k, _ = x.shape

buf = np.zeros_like(x, shape=(n, k, k + steps * (n - 1)))
buf[np.arange(n)[:, None, None], np.arange(k)[:, None], np.arange(k) + step * np.arange(n)[:, None, None]] = x

result = buf.sum(0)

我怀疑一个简单的循环会更快，因为你在这里索引浪费了很多空间。

这是一个快速的计时测试：

def index_emplace(x, steps):
    k, m, n = x.shape
    buf = np.zeros_like(x, shape=(k, m, n + steps * (k - 1)))
    r = np.arange(max(x.shape))
    a = r[:k][:, None, None]
    b = r[:m][:, None]
    c = r[:n]
    buf[a, b, c + step * a] = x
    return buf.sum(0)

def loop_emplace(x, steps):
    k, m, n = x.shape
    result = np.zeros_like(x, shape=(m, n + steps * (k - 1)))
    for i in range(k):
        result[:, step * i:step * i + m] += x[i]
    return result

x = np.random.randint(10, shape=(100, 100, 100))
steps = 37

%timeit index_emplace(x, steps)
107 ms ± 970 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit loop_emplace(x, steps)
2.26 ms ± 14.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

(index_emplace(x, steps) == loop_emplace(x, steps)).all()
True

如您所见，对于大型数组，循环可以快 50 倍。

推荐阅读