Question

如何在列表指定的每列中填充n个连续元素。 cumsum从下一行开始重置。例如，镜头= [2,3] ==＆gt;前2行[：2]的cumsum，接下来3行[2：5]

的cumsum

import numpy as np
lens = [3, 2]
a = np.array(
       [[ 1,  2],
        [ 1,  2],
        [ 1,  2],
        [ 1,  2],
        [ 1,  2]])

给

np.array(
       [[ 1,  2],
        [ 2,  4],
        [ 3,  6],
        [ 1,  2],
        [ 2,  4]])

试图避免循环

Answer 1

一个选项是拆分数组，cumsum然后将它们组合起来：

np.concatenate(list(map(lambda a: np.cumsum(a, axis=0), np.array_split(a, np.cumsum(lens)))))
#array([[1, 2],
#       [2, 4],
#       [3, 6],
#       [1, 2],
#       [2, 4]], dtype=int32)

没有拆分和组合的另一个选项是创建一个辅助数组，用于重置特定索引处的总和，如下所示：

idx = np.cumsum([0] + lens)[:-1]
aux = np.zeros_like(a)
aux[idx[1:], :] = -np.add.reduceat(a, idx)[:-1]
(a + aux).cumsum(0)

#array([[1, 2],
#       [2, 4],
#       [3, 6],
#       [1, 2],
#       [2, 4]], dtype=int32)

这两种方法速度大致相同：

def split_concat(a):
    return np.concatenate(list(map(lambda a: np.cumsum(a, axis=0), np.array_split(a, np.cumsum(lens)))))

def reset_sum(a):
    idx = np.cumsum([0] + lens)[:-1]
    aux = np.zeros_like(a)
    aux[idx[1:], :] = -np.add.reduceat(a, idx)[:-1]
    return (a + aux).cumsum(0)


lens = np.arange(1000)
a = np.ones((lens.sum(), 2))
(reset_sum(a) == split_concat(a)).all()
# True

%timeit split_concat(a)
# 12.8 ms ± 35.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit reset_sum(a)
# 13.6 ms ± 87.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

由列表

1 个答案: