如何从带列索引的NumPy矩阵中提取分组列？

嘿，这个需求很直观，我们可以利用NumPy的布尔索引轻松实现列的分组提取，下面是具体的步骤和代码：

思路解析 #

核心思路是：先找出索引向量里的所有唯一分组标识，然后对每个标识，筛选出原始矩阵中对应列，最后把这些子矩阵收集到列表中。

完整代码示例 #

import numpy as np

# 定义你的原始矩阵和索引向量
M = np.matrix([[1,2,3,4,5,6,7,8,9,10], 
               [11,12,13,14,15,16,17,18,19,20], 
               [21,22,23,24,25,26,27,28,29,30]])
index = np.array([1,1,2,2,2,2,3,4,4,4])

# 一行搞定分组提取（也可以拆成循环，看你习惯）
matrix_list = [M[:, index == val] for val in np.unique(index)]

# 可以把列表里的矩阵分别取出来验证
M1, M2, M3, M4 = matrix_list
print("M1:\n", M1)
print("\nM2:\n", M2)
print("\nM3:\n", M3)
print("\nM4:\n", M4)

代码解释 #

• np.unique(index)：提取索引向量里的所有唯一值，这里得到[1,2,3,4]，刚好对应我们需要的4个分组。
• index == val：生成一个布尔数组，标记哪些列属于当前分组（比如当val=1时，布尔数组前两位是True，其余是False）。
• M[:, index == val]：利用布尔索引提取矩阵中对应列，得到该分组的子矩阵。
• 列表推导式把所有子矩阵一次性收集到matrix_list里，简洁高效。

输出结果 #

运行代码后，你会得到完全符合需求的子矩阵：

M1:
[[ 1 2]
[11 12]
[21 22]]

M2:
[[ 3 4 5 6]
[13 14 15 16]
[23 24 25 26]]

M3:
[[ 7]
[17]
[27]]

M4:
[[ 8 9 10]
[18 19 20]
[28 29 30]]

内容的提问来源于stack exchange，提问作者Álvaro Méndez Civieta