python，对相同尺寸的numpy数组求和，不同形状的偏移和滚动结束

Question

我有 2 个 numpy 数组，想用偏移量对它们求和。 总和总是具有数组“a”的形状

a = np.array([1, 1, 1, 1, 1])
b = np.array([5, 6])
sumWithRoll(a, b, offset=1)
print(a)
>> [1, 6, 7, 1, 1]

此外，如果数组“b”很长或偏移量足够大，它应该翻转数组“a”的末尾：

a = np.array([1, 1, 1, 1, 1])
b = np.array([5, 6, 7, 8])
sumWithRoll(a, b, offset=3)
print(a)
>> [8, 9, 1, 6, 7]

我需要这个来合并两个循环播放的声音缓冲区，并希望有一个快速的解决方案，占用更少的内存。

编辑： 我有一个看起来很长的解决方案：

def sumWithRoll(buffer, indata, idx):
    buffLen = len(buffer)
    dataLen = len(indata)
    if dataLen > buffLen:
        indata = indata[0:buffLen]
        dataLen = buffLen
    idx = idx % buffLen
    idx2 = (idx + dataLen) % buffLen
    if idx2 <= idx:
        idx3 = buffLen - idx
        buffer[idx:buffLen] += indata[0:idx3]                   
        buffer[0:idx2] += indata[idx3:buffLen]
    else:
        buffer[idx:idx2] += indata[:]

希望有Pythonic的一两行解决方案

Answer 1

试试np.roll:

import numpy as np


def sum_with_roll(a, b, offset=0):
    e = np.zeros(a.shape)
    e[tuple(map(slice, b.shape))] = b
    return a + np.roll(e, offset)


a = np.array([1, 1, 1, 1, 1])
b = np.array([5, 6])
print(sum_with_roll(a, b, offset=1))

a = np.array([1, 1, 1, 1, 1])
b = np.array([5, 6, 7, 8])
print(sum_with_roll(a, b, offset=3))

输出：

[1. 6. 7. 1. 1.]
[8. 9. 1. 6. 7.]

---

对于list 输出：

def sum_with_roll(a, b, offset=0):
    e = np.zeros(a.shape)
    e[tuple(map(slice, b.shape))] = b
    return (a + np.roll(e, offset)).tolist()

[1.0, 6.0, 7.0, 1.0, 1.0]
[8.0, 9.0, 1.0, 6.0, 7.0]

---

对于int 类型输出：

def sum_with_roll(a, b, offset=0):
    e = np.zeros(a.shape)
    e[tuple(map(slice, b.shape))] = b
    return (a + np.roll(e, offset)).astype(int)

[1 6 7 1 1]
[8 9 1 6 7]

Answer 2

可以使用简单的模运算符来模拟滚动。

鉴于您想将a 和b 加在一起的一些索引，它基本上只是在做a[indices] += b。

a = np.array([1, 1, 1, 1, 1])
b = np.array([5, 6])
offset = 1

indices = np.arange(offset, len(b)+offset) % len(a)
a[indices] += b
# np.add.at(a, indices, b) - if you want to do the operation unbuffered

a 中的输出：

array([1, 6, 7, 1, 1])

a = np.array([1, 1, 1, 1, 1])
b = np.array([5, 6, 7, 8])
offset = 3

indices = np.arange(offset, len(b)+offset) % len(a)
a[indices] += b

a 中的输出：

array([8, 9, 1, 6, 7])

您也可以为此使用np.bincount：

a = np.array([1, 1, 1, 1, 1])
b = np.array([5, 6, 7, 8])
offset = 3

indices = np.arange(offset, len(b)+offset) % len(a)
np.bincount(indices, b)+a

输出：

array([8., 9., 1., 6., 7.])

很遗憾，无法为np.bincount 指定结果缓冲区，这会使+a 操作过时。

编辑

如果你有兴趣让它更快，我建议将就地添加分成 2 片：

a = np.array([1, 1, 1, 1, 1])
b = np.array([5, 6, 7, 8])
offset = 3

size = len(a) - offset
s1 = slice(offset, offset + size)
s2 = slice(None,size)

size = len(b) - size
s3 = slice(None, size)
s4 = slice(size, None)

a[s1] += b[s2]
a[s3] += b[s4]

输出：

array([8, 9, 1, 6, 7])