Python中快速构造稀疏块矩阵及三维场景扩展方法咨询

基础场景 #

希望在Python中构造形状为(N, 2N)的矩阵，当前实现使用循环：

import numpy as np 

N = 10 # 10,100,1000, whatever

some_vector = np.random.uniform(size=N) 

some_matrix = np.zeros((N, 2*N))

for i in range(N):
    some_matrix[i, 2*i] = 1
    some_matrix[i, 2*i + 1] = some_vector[i]

构造的矩阵大部分元素为0，仅第i行的2*i和2*i+1列有非零值，寻求无需循环的更快构造方式。

扩展场景 #

实际使用场景中，some_vector形状为(N,T)，需要构造形状为(N,2N,T)的矩阵，朴素实现如下：

N = 10  # 10,100,1000, whatever
T = 500 # or whatever

some_vector = np.random.uniform(size=(N,T)) 

some_matrix = np.zeros((N, 2*N,T))

for i in range(N):
    for t in range(T):
        some_matrix[i, 2*i,t] = 1
        some_matrix[i, 2*i + 1,t] = some_vector[i,t]

希望将基础场景的优化方案扩展到该场景。

基础场景优化 #

利用numpy的索引广播特性，直接通过数组索引批量赋值，完全避免循环：

import numpy as np

N = 10
some_vector = np.random.uniform(size=N)
some_matrix = np.zeros((N, 2*N))

# 生成行索引和对应的列索引
row_indices = np.arange(N)
col_indices_1 = 2 * row_indices
col_indices_2 = 2 * row_indices + 1

# 批量赋值
some_matrix[row_indices, col_indices_1] = 1
some_matrix[row_indices, col_indices_2] = some_vector

通过np.arange(N)获取所有行的索引，直接计算出每行需要赋值的两个列位置，一次性完成赋值操作，比循环效率高得多，尤其是当N很大时优势明显。

扩展场景优化 #

同样利用索引的维度扩展来匹配T维度，实现无循环赋值：

import numpy as np

N = 10
T = 500
some_vector = np.random.uniform(size=(N,T))
some_matrix = np.zeros((N, 2*N, T))

# 生成扩展维度后的索引，匹配T维度
row_indices = np.arange(N)[:, np.newaxis]  # 形状变为(N,1)，可与T维度广播
col_indices_1 = 2 * np.arange(N)[:, np.newaxis]
col_indices_2 = 2 * np.arange(N)[:, np.newaxis] + 1

# 批量赋值
some_matrix[row_indices, col_indices_1, :] = 1
some_matrix[row_indices, col_indices_2, :] = some_vector

将行索引和列索引扩展为(N,1)的形状，这样可以和第三个维度T进行广播匹配，一次性完成所有时间步t的赋值，彻底摆脱嵌套循环，大幅提升效率。

内容的提问来源于stack exchange，提问作者user1887919