Numpy: Efficiently create this Matrix (N,3) base values of another list and repeating them-CodePudding

How can I create the matrix

[[a, 0, 0],
 [0, a, 0],
 [0, 0, a],
 [b, 0, 0],
 [0, b, 0],
 [0, 0, b],
 ...]

from the vector

[a, b, ...]

efficiently?

There must be a better solution than

np.squeeze(np.reshape(np.tile(np.eye(3), (len(foo), 1, 1)) * np.expand_dims(foo, (1, 2)), (1, -1, 3)))

right?

CodePudding user response：

You can create a zero array in advance, and then quickly assign values by slicing:

def concated_diagonal(ar, col):
    ar = np.asarray(ar).ravel()
    size = ar.size
    ret = np.zeros((col * size, col), ar.dtype)
    for i in range(col):
        ret[i::col, i] = ar
    return ret

Test:

>>> concated_diagonal([1, 2, 3], 3)
array([[1, 0, 0],
       [0, 1, 0],
       [0, 0, 1],
       [2, 0, 0],
       [0, 2, 0],
       [0, 0, 2],
       [3, 0, 0],
       [0, 3, 0],
       [0, 0, 3]])

Note that because the number of columns you require is small, the impact of the relatively slow Python level for loop is acceptable:

>>> timeit(lambda: concated_diagonal(np.arange(1_000_000), 3), number=10)
0.1810098999994807
>>> # 18ms per loop

CodePudding user response：

Here is a solution by indexing:

a = [1,2,3]

N = 3
M = N*len(a)
out = np.zeros((M, N), dtype=int)
idx = np.arange(M)
out[idx, idx%N] = np.repeat(a, N)

output:

array([[1, 0, 0],
       [0, 1, 0],
       [0, 0, 1],
       [2, 0, 0],
       [0, 2, 0],
       [0, 0, 2],
       [3, 0, 0],
       [0, 3, 0],
       [0, 0, 3]])

intermediates:

idx
# array([0, 1, 2, 3, 4, 5, 6, 7, 8])

idx%N
# array([0, 1, 2, 0, 1, 2, 0, 1, 2])

np.repeat(a, N)
# array([1, 1, 1, 2, 2, 2, 3, 3, 3])

CodePudding user response：

You can use