pandas 找出相关性最强的两个变量

看过来

《pandas 教程》持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

在数据分析中经常要寻找数据之间的相关性。假设有这样一个数据集：

import pandas as pd

df = pd.DataFrame({
    'A': [1, 2, 4, 5, 6],
    'B': [2, 4, 6, 9, 10],
    'C': [2, 1, 7, 2, 1]
})
df
'''
   A   B  C
0  1   2  2
1  2   4  1
2  4   6  7
3  5   9  2
4  6  10  1
'''

需要找出 A、B、C 三列中相关性最强的两列。我们知道使用 df.corr() 可以得到这三列的相关性矩阵。相关性系数从 –1 到 1 表示相关程度。代码如下：

# 相关性
df.corr()
'''
          A         B         C
A  1.000000  0.987069  0.057639
B  0.987069  1.000000 -0.077381
C  0.057639 -0.077381  1.000000
'''

可以看到，从左上到右下对角线上的值全为 1，这些值是列和自己的相关性，要去掉。现在要找到去掉这些值后相关性系数最大的两列。利用 stack() 堆叠数据，转为一列，行和列上的轴标签形成两层索引，然后对值从大到小排序：

(
    df.corr()  # 相关性矩阵
    .stack()  # 堆叠，转为一列
    .sort_values(ascending=False)  # 排序，最大值在前
)
'''
A  A    1.000000
B  B    1.000000
C  C    1.000000
A  B    0.987069
B  A    0.987069
A  C    0.057639
C  A    0.057639
B  C   -0.077381
C  B   -0.077381
dtype: float64
'''

根据前面的分析，筛选掉值为 1 的数据：

(
    df.corr()  # 相关性矩阵
    .stack()  # 堆叠，转为一列
    .sort_values(ascending=False)  # 排序，最大值在前
    .loc[lambda x:x < 1]  # 去掉值为1的数据
)
'''
A  B    0.987069
B  A    0.987069
A  C    0.057639
C  A    0.057639
B  C   -0.077381
C  B   -0.077381
dtype: float64
'''

我们看到最大的值为 0.987069，它的标签是 B 和 A，需求是希望知道标签而不是最大值是多少，怎么实现呢？可以用 idxmax()，它可以得到最大值的索引：

(
    df.corr()  # 相关性矩阵
    .stack()  # 堆叠，转为一列
    .sort_values(ascending=False)  # 排序，最大值在前
    .loc[lambda x:x < 1]  # 去掉值为1的数据
    .idxmax()  # 最大值的标签
)
# ('A', 'B')

这样就得到了一个元组，元组里的两个元素就是相关性最强的两列的标签。

（完）

注：此案例收录在《深入浅出Pandas：利用Python进行数据处理与分析》17.2.4 小节。

pandas 找出相关性最强的两个变量

相关内容