如何根据N×M数组行的取值组合提取对应元素

示例数据定义 #

先定义题目中的两个Numpy数组：

import numpy as np

# 浮点数值数组 (2, 10)
float_arr = np.array([
    [0.1, 0.02, 0.2, 0.3, 0.013, 0.7, 0.7, 0.11, 0.18, 0.6],
    [0.23, 0.02, 0.1, 0.1, 0.011, 0.3, 0.4, 0.4, 0.4, 0.5]
])

# 二进制组合数组 (2, 10)
binary_arr = np.array([
    [0, 1, 0, 0, 0, 0, 1, 1, 0, 1],
    [1, 0, 1, 0, 0, 0, 0, 1, 1, 0]
])

实现步骤 #

1. 转置二进制数组：将原数组的列（每个位置的组合）转为行，方便Numpy按行识别唯一组合

   # 转置后形状为(10, 2)，每一行对应原数组的一个位置组合
   binary_combs = binary_arr.T
   

2. 提取唯一组合：用np.unique获取所有不重复的二进制组合

   unique_combs = np.unique(binary_combs, axis=0)
   

3. 遍历组合提取对应元素：对每个唯一组合，生成掩码筛选匹配的位置，再提取浮点数组的对应元素

   for comb in unique_combs:
       # 生成掩码：标记所有匹配当前组合的位置
       mask = np.all(binary_combs == comb, axis=1)
       # 提取浮点数组中对应位置的元素，结果形状为(2, k)，k是匹配的位置数量
       extracted_elements = float_arr[:, mask]
       
       print(f"组合 {tuple(comb)} 对应的元素：")
       print(extracted_elements)
       print("---")
   

运行结果 #

执行代码后会输出每个唯一组合对应的元素：

组合 (0, 0) 对应的元素：
[[0.3   0.013]
 [0.1   0.011]]
---
组合 (0, 1) 对应的元素：
[[0.1  0.2  0.18]
 [0.23 0.1  0.4 ]]
---
组合 (1, 0) 对应的元素：
[[0.02 0.7  0.6 ]
 [0.02 0.4  0.5 ]]
---
组合 (1, 1) 对应的元素：
[[0.11]
 [0.4 ]]
---

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