对比多个序列的相同元素占比并绘制热力图-Python

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

'''
对比两个数组前k个元素的相同元素占比
'''
def getArraySame(array1, array2, k):
    length = len(array1)
    if k > length:
        k = length
    # 统计两个数组相同元素个数
    count = len(set(array1[0:k]) & set(array2[0:k]))
    # 除以总元素数，转为占比
    return count / k


if __name__ == '__main__':
    '''
    数据准备
    '''
    methods = pd.read_csv('data.csv')  # 多种方法的排序结果
    # 文件中第一列为排名，因此方法名从索引1至n
    methods_name = np.array(methods.columns)[1: len(methods.columns)]  
    m = len(methods_name) # 共多少个方法

    '''
    计算
    '''
    a = np.zeros((m, m)) # 初始化混淆矩阵
    k = int(input("请输入您想要对比的排名情况(截止至k):"))  # 前k个元素
    for i in range(m):
        for j in range(m):
            # 调用对比方法，比较2个序列的相同元素占比
            a[i][j] = getArraySame(methods[methods_name[i]].tolist(), methods[methods_name[j]].tolist(), k)
    
    '''
    存储至文件
    '''
    df = pd.DataFrame(a, index=methods_name, columns=methods_name)  
    print(df)
    df.to_csv('Top-' + str(k) + ' 序列 相同元素占比.csv', index=False)
    print("已输出到文件！")

    '''
    画热力图
    '''
    sns.heatmap(df, cmap="YlGnBu_r", annot=True, annot_kws={'size': 9, 'weight': 'bold', 'color': 'black'})
    plt.title(" Top-" + str(k))
    plt.show()