按非零元素数量水平拆分CSR矩阵的实现异常问题

问题原因 #

你写的代码确实是因为submatrices.append(row)和逐行递增currentRow导致生成单行子矩阵的。当前逻辑每次循环只把单独一行加入结果列表，而没有将连续的行合并成一个完整的子矩阵。

修正思路 #

要实现按非零元素数量分配行的目标，需要：

• 记录每个子矩阵的起始行位置
• 累计当前子矩阵的非零元素数量，直到接近目标值时，一次性切分CSR矩阵的连续行区间作为子矩阵
• 处理最后一个分区，确保所有剩余行都被包含

修正后的代码 #

from scipy import sparse

def horizontalSplit(matrix, numPartitions):
    csr = sparse.csr_matrix(matrix)
    submatrices = []
    row_nonzeros = csr.getnnz(axis=1)
    total_nonzeros = sum(row_nonzeros)
    target_nonzeros = total_nonzeros // numPartitions
    
    current_row = 0
    start_row = 0
    current_nonzeros = 0
    
    for i in range(numPartitions):
        # 累计行直到达到目标非零元素，或者到最后一行
        while current_row < csr.shape[0] and current_nonzeros < target_nonzeros:
            current_nonzeros += row_nonzeros[current_row]
            current_row += 1
        # 处理最后一个分区，直接取剩余所有行
        if i == numPartitions - 1:
            submat = csr[start_row:]
        else:
            submat = csr[start_row:current_row]
        submatrices.append(submat)
        # 重置起始行和计数
        start_row = current_row
        current_nonzeros = 0
    
    return submatrices

# 测试用例
adjacency_matrix = [
    [1, 1, 1, 1, 0, 0, 0, 0],
    [1, 0, 1, 0, 0, 0, 0, 0],
    [1, 1, 0, 1, 0, 0, 0, 0],
    [1, 0, 1, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 1, 1, 1],
    [0, 0, 0, 0, 1, 0, 1, 0],
    [0, 0, 0, 0, 1, 1, 0, 1],
    [0, 0, 0, 0, 1, 0, 1, 0]
]
csr_matrix = sparse.csr_matrix(adjacency_matrix)
subs = horizontalSplit(csr_matrix, 4)

# 打印每个子矩阵的非零元素数量和内容
for idx, sub in enumerate(subs):
    print(f"=== 子矩阵 {idx+1} ===")
    print(f"非零元素总数: {sub.getnnz()}")
    print(sub.toarray())
    print("------------------")

运行结果解释 #

对于测试的8x8矩阵，总非零元素为21，每个分区目标为5：

1. 子矩阵1：行0（4个非零）+行1（2个）=6，接近目标值5，作为第一个子矩阵
2. 子矩阵2：行2（3个）+行3（2个）=5，刚好达标
3. 子矩阵3：行4（3个）+行5（2个）=5，刚好达标
4. 子矩阵4：剩余行6（3个）+行7（2个）=5，总计21个非零元素分配完毕

输出示例：

=== 子矩阵 1 ===
非零元素总数: 6
[[1 1 1 1 0 0 0 0]
 [1 0 1 0 0 0 0 0]]
------------------
=== 子矩阵 2 ===
非零元素总数: 5
[[1 1 0 1 0 0 0 0]
 [1 0 1 0 0 0 0 0]]
------------------
=== 子矩阵 3 ===
非零元素总数: 5
[[0 0 0 0 0 1 1 1]
 [0 0 0 0 1 0 1 0]]
------------------
=== 子矩阵 4 ===
非零元素总数: 5
[[0 0 0 0 1 1 0 1]
 [0 0 0 0 1 0 1 0]]
------------------

额外说明 #

如果需要同时满足非零元素均衡和行数对齐64的要求，可以在调整子矩阵行范围时，最后将起始/结束行对齐到64的倍数（比如结束行向上取整到最近的64倍数），但需要注意最后一个分区可能会包含更多行，避免非零元素差异过大。

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