无需迭代在Python中为贷款数据生成相关矩阵的优化方案

需求明确 #

给定包含country（国家）、sector（行业）、loan（贷款标识）的数据集，需生成相关系数矩阵，规则优先级从高到低为：

• 同一贷款（loan字段值相同）：相关系数为1
• 同行业且同国家：相关系数为35%（0.35）
• 仅同国家 或 仅同行业：相关系数为25%（0.25）
• 其余情况：相关系数为5%（0.05）

优化思路：向量化替代嵌套循环 #

Python嵌套循环在数据量大时效率极低，利用NumPy的广播机制和outer方法生成布尔矩阵，可实现完全无迭代的高效计算，逻辑清晰且性能提升显著。

实现代码 #

import numpy as np
import pandas as pd

# 从数据集中提取核心字段的NumPy数组
loan_arr = df['loan'].values
sector_arr = df['sector'].values
country_arr = df['country'].values
n_rows = len(loan_arr)

# 1. 初始化矩阵为默认值（其余情况：5%）
corr_matrix = np.full((n_rows, n_rows), 0.05)

# 2. 生成各条件的布尔矩阵
same_loan = np.equal.outer(loan_arr, loan_arr)
same_sector = np.equal.outer(sector_arr, sector_arr)
same_country = np.equal.outer(country_arr, country_arr)

# 3. 按优先级依次覆盖数值
# 仅同国家/仅同行业的情况
only_same_sector_or_country = np.logical_and(
    np.logical_or(same_sector, same_country),
    ~np.logical_and(same_sector, same_country)
)
corr_matrix[only_same_sector_or_country] = 0.25

# 同行业且同国家的情况
corr_matrix[np.logical_and(same_sector, same_country)] = 0.35

# 同一贷款的情况（最高优先级）
corr_matrix[same_loan] = 1

代码说明 #

• np.equal.outer(a, b)：生成二维布尔矩阵，矩阵中[i,j]位置的值为a[i] == b[j]，批量完成元素对比，无需循环。
• 按优先级从低到高初始化赋值：先设默认值，再用低优先级条件覆盖，最后用最高优先级条件覆盖，确保逻辑正确。

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