将元素插入numpy数组并获取所有滚动排列

Question

有没有更好的方法将数组中的元素逐个插入
到所有可能的位置(n + 1个位置).

例如,插入[1]到[6 7 8 9]应该产生:

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

所以,如果我A = [1 2 3]逐个插入B = [6 7 8 9]它应该产生:

[1 6 7 8 9]
[9 1 6 7 8]
[8 9 1 6 7]
[7 8 9 1 6]
[6 7 8 9 1]
--------------------
[2 6 7 8 9]
[9 2 6 7 8]
[8 9 2 6 7]
[7 8 9 2 6]
[6 7 8 9 2]
--------------------
[3 6 7 8 9]
[9 3 6 7 8]
[8 9 3 6 7]
[7 8 9 3 6]
[6 7 8 9 3]
--------------------

目前我使用numpy.roll这样的:

import numpy as np
import timeit

A = np.array([1, 2, 3, 4, 5])
B = np.array([6, 7, 8, 9])

def inject_one(Ad, Bd):
    for i, _ in enumerate(Ad):
        C = np.append(Ad[i], Bd)
        for _ in range(len(C) - 1):
            C = np.roll(C, 1)

t = timeit.Timer(lambda: inject_one(A, B))
print("{:.3f}secs for 1000 iterations".format(t.timeit(number=1000))) 
# > 0.160 secs

Answer 1

您在这里要求的是Toeplitz Matrix，它是：

一个矩阵，其中从左到右的每个降对角线都是常数

幸运的是，scipy有一个易于使用的实现：

from scipy.linalg import toeplitz

def magic_toeplitz(arr, to_add):
    return toeplitz(np.hstack([to_add, arr[::-1]]), np.hstack([to_add, arr]))

a = [6,7,8,9]
add = [1]
magic_toeplitz(a, add)

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

缩放该解决方案的矢量化方法：

A = np.array([1, 2, 3, 4, 5])
B = np.array([6, 7, 8, 9])

out = toeplitz(np.hstack([[np.nan], B[::-1]]), np.hstack([np.nan, B]))
out = np.tile(out, (len(A), 1, 1))
m = np.ma.array(out, mask=np.isnan(out))
vals = np.repeat(A, (B.shape[0] + 1)**2).reshape(out.shape)
print(m.filled(vals))

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

       [[2, 6, 7, 8, 9],
        [9, 2, 6, 7, 8],
        [8, 9, 2, 6, 7],
        [7, 8, 9, 2, 6],
        [6, 7, 8, 9, 2]],

       [[3, 6, 7, 8, 9],
        [9, 3, 6, 7, 8],
        [8, 9, 3, 6, 7],
        [7, 8, 9, 3, 6],
        [6, 7, 8, 9, 3]],

       [[4, 6, 7, 8, 9],
        [9, 4, 6, 7, 8],
        [8, 9, 4, 6, 7],
        [7, 8, 9, 4, 6],
        [6, 7, 8, 9, 4]],

       [[5, 6, 7, 8, 9],
        [9, 5, 6, 7, 8],
        [8, 9, 5, 6, 7],
        [7, 8, 9, 5, 6],
        [6, 7, 8, 9, 5]]])