在pandas数据框中返回相关列的分组

10 投票
2 回答
6650 浏览
提问于 2025-04-18 08:19

我在一个 pandas DataFrame 上运行了一个相关性矩阵:

df=pd.DataFrame( {'one':[0.1, .32, .2, 0.4, 0.8], 'two':[.23, .18, .56, .61, .12], 'three':[.9, .3, .6, .5, .3], 'four':[.34, .75, .91, .19, .21], 'zive': [0.1, .32, .2, 0.4, 0.8], 'six':[.9, .3, .6, .5, .3], 'drive':[.9, .3, .6, .5, .3]})

corrMatrix=df.corr()
corrMatrix
           drive  four   one   six  three   two  zive
drive       1.00 -0.04 -0.75  1.00   1.00  0.24 -0.75
four       -0.04  1.00 -0.49 -0.04  -0.04  0.16 -0.49
one        -0.75 -0.49  1.00 -0.75  -0.75 -0.35  1.00
six         1.00 -0.04 -0.75  1.00   1.00  0.24 -0.75
three       1.00 -0.04 -0.75  1.00   1.00  0.24 -0.75
two         0.24  0.16 -0.35  0.24   0.24  1.00 -0.35
zive       -0.75 -0.49  1.00 -0.75  -0.75 -0.35  1.00

现在,我想写一些代码,返回那些完全相关的列(也就是说,相关性等于1),并把它们分组。

理想情况下,我想要这样的结果: [['zive', 'one'], ['three', 'six', 'drive']]

我写了下面的代码,结果是 ['drive', 'one', 'six', 'three', 'zive'],但正如你所看到的,这只是一些有完美相关性的列的集合——它们没有被分成明显的组,和它们完美相关的列在一起。

correlatedCols=[]
for col in corrMatrix:
    data=corrMatrix[col][corrMatrix[col]==1]
    if len(data)>1:
        correlatedCols.append(data.name)

correlatedCols  
['drive','one', 'six', 'three', 'zive']

编辑: 根据 @Karl D. 提供的建议,我得到了这个结果:

cor = df.corr()
cor.loc[:,:] =  np.tril(cor.values, k=-1)
cor = cor.stack()
cor[cor ==1]
six    drive   1.00
three  drive   1.00
       six     1.00
zive   one     1.00

..但这并不是我想要的——因为 [six, drive] 不是一个完整的分组——它缺少了 'three'

2 个回答

10

你可以尝试下面这样的做法:

>>> cor = df.corr()
>>> cor.loc[:,:] =  np.tril(cor, k=-1)
>>> cor = cor.stack()
>>> cor[cor > 0.9999]

three  six    1
zive   one    1

如果想让结果更接近你期望的输出,可以这样做:

>>> cor[cor > 0.9999].to_dict().keys()

[('zive', 'one'), ('three', 'six')]

解释一下。首先,我创建了一个只包含下三角部分的协方差矩阵,并且不包括对角线上的元素(使用numpy的 tril):

>>> cor.loc[:,:] =  np.tril(cor.values, k=-1)

           four       one       six     three       two  zive
four   0.000000 -0.000000 -0.000000 -0.000000  0.000000    -0
one   -0.489177  0.000000 -0.000000 -0.000000 -0.000000     0
six   -0.039607 -0.747365  0.000000  0.000000  0.000000    -0
three -0.039607 -0.747365  1.000000  0.000000  0.000000    -0
two    0.159583 -0.351531  0.238102  0.238102  0.000000    -0
zive  -0.489177  1.000000 -0.747365 -0.747365 -0.351531     0

然后我把数据框(dataframe)堆叠起来:

>>> cor = cor.stack()

four   four     0.000000
       one     -0.000000
       six     -0.000000
       three   -0.000000
       two      0.000000
       zive    -0.000000
one    four    -0.489177
       one      0.000000
       six     -0.000000
       three   -0.000000
       two     -0.000000
       zive     0.000000
six    four    -0.039607
       one     -0.747365
       six      0.000000
       three    0.000000
       two      0.000000
       zive    -0.000000
three  four    -0.039607
       one     -0.747365
       six      1.000000
       three    0.000000
       two      0.000000
       zive    -0.000000
two    four     0.159583
       one     -0.351531
       six      0.238102
       three    0.238102
       two      0.000000
       zive    -0.000000
zive   four    -0.489177
       one      1.000000
       six     -0.747365
       three   -0.747365
       two     -0.351531
       zive     0.000000

接着,我可以直接获取那些值为一的行。

编辑:我觉得这样可以得到你想要的格式,但看起来不是很优雅:

>>> from itertools import chain

>>> cor.loc[:,:] =  np.tril(cor, k=-1)
>>> cor = cor.stack()
>>> ones = cor[cor > 0.999].reset_index().loc[:,['level_0','level_1']]
>>> ones = ones.query('level_0 not in level_1')
>>> ones.groupby('level_0').agg(lambda x: set(chain(x.level_0,x.level_1))).values

[[set(['six', 'drive', 'three'])]
 [set(['zive', 'one'])]]
4

这里有一个简单的方法:

df=pd.DataFrame( {'one':[0.1, .32, .2, 0.4, 0.8], 'two':[.23, .18, .56, .61, .12], 'three':[.9, .3, .6, .5, .3], 'four':[.34, .75, .91, .19, .21], 'zive': [0.1, .32, .2, 0.4, 0.8], 'six':[.9, .3, .6, .5, .3], 'drive':[.9, .3, .6, .5, .3]})

corrMatrix=df.corr()

corrMatrix.loc[:,:] =  np.tril(corrMatrix, k=-1) # borrowed from Karl D's answer

already_in = set()
result = []
for col in corrMatrix:
    perfect_corr = corrMatrix[col][corrMatrix[col] == 1].index.tolist()
    if perfect_corr and col not in already_in:
        already_in.update(set(perfect_corr))
        perfect_corr.append(col)
        result.append(perfect_corr)

结果:

>>> result
[['six', 'three', 'drive'], ['zive', 'one']]

撰写回答