2. COO 到 CSR 的转换实现与正确性验证

问题拆解：COO 的三元组如何变成 CSR 的行结构

COO 用三个等长数组 row、col、val 逐条记录非零元，天然允许重复索引且无序；CSR 则要求按行分组、行内列索引升序，并用 indptr 标记每行在 indices/data 中的起止位置。转换的核心不是简单搬运，而是完成三件事：重复项合并（求和）、按行稳定分组、行内按列排序，最后生成长度为 n_rows+1 的行指针。

转换算法：计数排序两遍扫描

最稳妥的实现是“计数—前缀和—填充”三段式，时间复杂度 O(nnz + n_rows)，且天然得到行内按原顺序的稳定排列。若输入未排序，需在填充后对每行做一次列索引排序（或先对三元组整体按 (row,col) 排序再合并）。下面给出完整可运行实现，包含重复合并与排序。

import numpy as np

def coo_to_csr(row, col, val, n_rows, n_cols, sort_cols=True):
    row = np.asarray(row, dtype=np.int64)
    col = np.asarray(col, dtype=np.int64)
    val = np.asarray(val, dtype=np.float64)
    nnz = row.size

    # 1) 合并重复索引：按 (row, col) 排序后相邻累加
    order = np.lexsort((col, row))          # 先按 row，再按 col
    row_s, col_s, val_s = row[order], col[order], val[order]

    keep_r, keep_c, keep_v = [], [], []
    for i in range(nnz):
        if keep_r and row_s[i] == keep_r[-1] and col_s[i] == keep_c[-1]:
            keep_v[-1] += val_s[i]          # 重复项求和
        else:
            keep_r.append(row_s[i]); keep_c.append(col_s[i]); keep_v.append(val_s[i])
    keep_r = np.asarray(keep_r, dtype=np.int64)
    keep_c = np.asarray(keep_c, dtype=np.int64)
    keep_v = np.asarray(keep_v, dtype=np.float64)

    # 2) 计数：每行非零元个数
    counts = np.bincount(keep_r, minlength=n_rows)

    # 3) 前缀和得到 indptr
    indptr = np.zeros(n_rows + 1, dtype=np.int64)
    indptr[1:] = np.cumsum(counts)

    # 4) 填充 indices / data（keep_r 已按行升序，故直接顺序写入）
    indices = keep_c.copy()
    data = keep_v.copy()

    # 5) 行内列索引排序（若 sort_cols 且未保证有序）
    if sort_cols:
        for r in range(n_rows):
            s, e = indptr[r], indptr[r + 1]
            if e - s > 1:
                seg = np.argsort(indices[s:e], kind="stable")
                indices[s:e] = indices[s:e][seg]
                data[s:e] = data[s:e][seg]

    return indptr, indices, data

# ---- 小规模稠密矩阵对照验证 ----
row = [2, 0, 1, 0, 2, 1]
col = [1, 0, 2, 0, 1, 2]
val = [5, 1, 3, 4, 7, 6]
n_rows, n_cols = 3, 3

indptr, indices, data = coo_to_csr(row, col, val, n_rows, n_cols)
print("indptr =", indptr)   # [0 1 3 4]
print("indices=", indices)  # [0 2 2 1]
print("data   =", data)     # [5. 3. 6. 12.]

# 稠密对照
dense = np.zeros((n_rows, n_cols))
for r, c, v in zip(row, col, val):
    dense[r, c] += v
print("dense =\n", dense)

# 由 CSR 还原稠密并比对
recon = np.zeros((n_rows, n_cols))
for r in range(n_rows):
    for k in range(indptr[r], indptr[r + 1]):
        recon[r, indices[k]] = data[k]
assert np.allclose(recon, dense), "CSR 与稠密矩阵不一致"
print("验证通过：CSR 与稠密矩阵逐元素一致")

正确性验证的三条硬指标

• 结构不变量：indptr[0]==0、indptr[-1]==len(data)==len(indices)、indptr 单调不减；每行 indices[indptr[r]:indptr[r+1]] 严格升序且无重复列。
• 数值等价：用 CSR 还原稠密矩阵，与 COO 直接累加得到的稠密矩阵做 np.allclose 比对，容差取 1e-12 量级。
• 边界覆盖：空行（counts 为 0）、全零矩阵、单行/单列、含负值或浮点重复项，均应通过上述两项检查。