在numpy中重复，但大小不一

Question

我有一个数组，它是不同块的串联：

a = np.array([0, 1, 2, 10, 11, 20, 21, 22, 23])
#             >     <  >    <  >            <
chunks = np.array([3, 2, 4])
repeats = np.array([1, 3, 2])

在上面的例子中，从新的十年开始的每个部分都是一个单独的“块”，我想重复一下。每个组块的大小和重复次数都是已知的。我无法执行后跟kron或repeat的重塑，因为块的大小不同

我想要的结果是

np.array([0, 1, 2, 10, 11, 10, 11, 10, 11, 20, 21, 22, 23, 20, 21, 22, 23])
# repeats:>  1  <  >         3          <  >              2             <

这在循环中很容易做到：

in_offset = np.r_[0, np.cumsum(chunks[:-1])]
out_offset = np.r_[0, np.cumsum(chunks[:-1] * repeats[:-1])]
output = np.zeros((chunks * repeats).sum(), dtype=a.dtype)
for c in range(len(chunks)):
    for r in range(repeats[c]):
        for i in range(chunks[c]):
            output[out_offset[c] + r * chunks[c] + i] = a[in_offset[c] + i]

这将导致以下矢量化：

regions = chunks * repeats
index = np.arange(regions.sum())

segments = np.repeat(chunks, repeats)
resets = np.cumsum(segments[:-1])
offsets = np.zeros_like(index)
offsets[resets] = segments[:-1]
offsets[np.cumsum(regions[:-1])] -= chunks[:-1]

index -= np.cumsum(offsets)

output = a[index]

有没有更有效的方法来矢量化这个问题？我们很清楚，我不是在要求代码审查。我很满意这些函数调用是如何协同工作的。我想知道是否有一种完全不同的（更有效的）函数调用组合可以用来实现相同的结果

这个问题是受{a1}到{a2}启发提出的

Answer 1

对于那些作为范围数组的块，我们可以直接处理输入数组，从而避免最后的索引步骤，这应该会有所改进-

# https://stackoverflow.com/a/47126435/ @Divakar
def create_ranges(starts, ends, l):
    clens = l.cumsum()
    ids = np.ones(clens[-1],dtype=int)
    ids[0] = starts[0]
    ids[clens[:-1]] = starts[1:] - ends[:-1]+1
    out = ids.cumsum()
    return out

s = np.r_[0,chunks.cumsum()]
starts = a[np.repeat(s[:-1],repeats)]
l = np.repeat(chunks, repeats)
ends = starts+l
out = create_ranges(starts, ends, l)

Answer 2

一种比另一种答案更能完成任务的方法是：

result = np.concatenate([ np.tile(tbl, rpt) for tbl, rpt in
    zip(np.split(a, np.cumsum(chunks[:-1])), repeats) ])

结果是：

array([ 0,  1,  2, 10, 11, 10, 11, 10, 11, 20, 21, 22, 23, 20, 21, 22, 23])

Answer 3

与另一个答案相比，解决这个问题的一个更为“numpythonic”的方法是-

np.concatenate(np.repeat(np.split(a, np.cumsum(chunks))[:-1], repeats))

array([ 0,  1,  2, 10, 11, 10, 11, 10, 11, 20, 21, 22, 23, 20, 21, 22, 23])

请注意，没有显式for循环

（np.split有一个隐式循环，正如@Divakar所指出的）

---

编辑：基准测试（MacBookPro 13）－

正如@Mad Physicast在他的帖子中指出的，Divakar的解决方案对于更大的阵列、块和重复来说可以更好地扩展