如何用Pandas DataFrame的值覆盖Scipy稀疏COO矩阵

我用Scipy构建了一个稀疏COO矩阵，示例数据如下：

(100981, 0)   0.015
  (100982, 0)   0.016
  (100983, 0)   0.016
  (100981, 1)   0.023
  (100982, 1)   0.024
  (100983, 1)   0.025
  (100981, 2)   0.019
  (100982, 2)   0.018
  (100983, 2)   0.017

同时有一个Pandas DataFrame，结构如下：

site_0   site_1   site_2    site
0  0.0140  0.02313  0.01020  100981
1  0.0130  0.02200  0.01370  100982
2  0.0162  0.02510  0.01479  100983

原先用Pandas apply逐行更新矩阵的代码：

# where v_coo is my matrix defined above
# where df is my dataframe defined above
def override_sites(row):
    v_coo[row['site'], 0] = row['site_0']
    v_coo[row['site'], 1] = row['site_1']
    v_coo[row['site'], 2] = row['site_2']
    return v_coo
df.apply(override_sites, axis=1)

问题说明 #

apply逐行更新效率极低，而且COO矩阵本身是不可变类型，每次赋值都会触发内部类型转换，性能损耗很大。下面是两种用NumPy向量化实现的高效方案：

方案1：转CSR矩阵后批量更新 #

CSR矩阵支持基于坐标的批量赋值，适合保留原矩阵大部分条目、仅更新小部分的场景：

import numpy as np
from scipy.sparse import csr_matrix

# 1. 将原COO矩阵转为CSR（支持高效赋值操作）
v_csr = v_coo.tocsr()

# 2. 构造向量化的行、列索引和对应值
rows = df['site'].repeat(3).values  # 每个site对应3列，重复3次
cols = np.tile([0, 1, 2], len(df))  # 列索引循环生成[0,1,2]
values = df[['site_0', 'site_1', 'site_2']].values.flatten()  # 扁平化所有待更新值

# 3. 批量更新矩阵
v_csr[rows, cols] = values

# 若需转回COO格式
v_coo_updated = v_csr.tocoo()

方案2：直接构造新的COO矩阵 #

利用COO矩阵基于row/col/data数组存储的特性，直接合并原矩阵条目和更新条目，去重后生成最终矩阵，适合大规模数据场景：

import numpy as np
from scipy.sparse import coo_matrix

# 1. 提取原COO矩阵的核心数据
orig_rows = v_coo.row
orig_cols = v_coo.col
orig_data = v_coo.data

# 2. 构造DataFrame对应的更新条目
update_rows = df['site'].repeat(3).values
update_cols = np.tile([0, 1, 2], len(df))
update_data = df[['site_0', 'site_1', 'site_2']].values.flatten()

# 3. 合并原条目与更新条目
all_rows = np.concatenate([orig_rows, update_rows])
all_cols = np.concatenate([orig_cols, update_cols])
all_data = np.concatenate([orig_data, update_data])

# 4. 去重：保留每个(row,col)对的最后一个值（即更新后的值）
coord_pairs = np.stack([all_rows, all_cols], axis=1)
_, unique_idx = np.unique(coord_pairs, axis=0, return_index=True)
unique_idx = np.sort(unique_idx)  # 可选：保持索引顺序，不影响矩阵内容

# 5. 生成更新后的COO矩阵
v_coo_updated = coo_matrix(
    (all_data[unique_idx], (all_rows[unique_idx], all_cols[unique_idx])),
    shape=v_coo.shape
)

为什么这两种方法更高效 #

所有操作都是基于NumPy数组的批量运算，完全避免了逐行循环，效率比apply高几个数量级，数据量越大优势越明显。

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