如何用Python将N×N列联矩阵转换为一维向量形式？

问题背景 #

我们需要将N×N列联矩阵转换为两个一维向量，以适配sklearn的调整互信息（AMI）函数——因为该函数不支持直接输入列联矩阵。

举个例子：

2×2列联矩阵：

Cj
    2  1
Ci
    1  0

转换后得到两个一维向量：

[0 0 0 1]
[0 0 1 0]

该列联矩阵代表两种各含2个簇的聚类算法结果：第一行表示Ci的簇1有3个数据点，簇2有1个数据点；Cj的簇A有3个数据点，簇B有1个数据点，两种算法在N=4个数据点中有2个达成一致。

另一个示例：

列联矩阵：

2 0 0
0 1 0
0 0 1

转换得到的两个一维向量：

0 0 1 2
0 0 1 2

转换逻辑如下：

V1
V2
For i 行索引 
  For j 列索引
     向V1添加contingency_ij个值为i的元素，向V2添加contingency_ij个值为j的元素

高效实现代码 #

利用numpy的向量化操作可以避免低效的Python嵌套循环，实现高效转换：

import numpy as np

def contingency_to_vectors(contingency):
    # 生成行列索引网格
    i_indices, j_indices = np.indices(contingency.shape)
    # 展平索引和计数
    i_flat = i_indices.flatten()
    j_flat = j_indices.flatten()
    counts = contingency.flatten()
    # 按计数重复索引
    v1 = np.repeat(i_flat, counts)
    v2 = np.repeat(j_flat, counts)
    return v1, v2

# 测试第一个示例
contingency1 = np.array([[2, 1], [1, 0]])
v1, v2 = contingency_to_vectors(contingency1)
print(v1)  # 输出: [0 0 0 1]
print(v2)  # 输出: [0 0 1 0]

# 测试第二个示例
contingency2 = np.array([[2, 0, 0], [0, 1, 0], [0, 0, 1]])
v1, v2 = contingency_to_vectors(contingency2)
print(v1)  # 输出: [0 0 1 2]
print(v2)  # 输出: [0 0 1 2]

代码说明 #

• np.indices快速生成与列联矩阵同形状的索引网格，一次性获取所有位置的行、列索引
• 展平操作将二维索引转换为一维数组，方便后续批量处理
• np.repeat基于每个位置的计数重复对应索引，numpy的底层C实现比纯Python循环效率高得多，尤其适合处理大尺寸列联矩阵

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