python - 编写一个以正类为输入的混淆矩阵函数

Question

我想创建一个混淆矩阵，而不依赖于任何包。我有两个列表（预测值和实际值），我想将它们与正类的指标一起输入到函数中。

例如，什么时候1是正类：

predicted_lst = [1, 0, 1, 0, 0]
actual_lst = [1, 0, 0, 1, 1]

我的功能目前看起来像这样，但效率低下：

def confusion_matrix(predicted, actual, pos_class):
    
    TP = 0
    TN = 0 
    FP = 0
    FN = 0
    
    for i in range(len(actual)):
        if actual[i] == pos_class and predicted[i] == pos_class:
            TP +=1
        elif actual[i] == pos_class and predicted[i] != pos_class:
            FN +=1
        elif actual[i] != pos_class and predicted[i] == pos_class:
            FP +=1
        else:
            TN +=1
    return TP, FP, TN, FN

我的问题是，有没有更有效的方法来编写这段代码？我看到了这些帖子，但它们没有像我希望的那样将正类作为函数输入。我也不想使用任何包（包括numpy）

标签： pythonmachine-learningconfusion-matrix

Answer 1

您可以假设它1是正类来编写函数，并pos_class通过简单地更改返回值的顺序来适应参数，如果0是正类的话。

FP在循环内部，您可以放弃and的增量计算FN，因为这些可以从循环外部的真实值中得出：

def confusion_matrix(predicted, actual, pos_class):
    
    TP = 0
    TN = 0 
    
    for pred, act in zip(predicted, actual):
        if pred == act:
            if act == 0:
                TN += 1
            else:
                TP += 1
    
    positive = sum(predicted)
    negative = len(predicted) - positive        
    FP = positive - TP 
    FN = negative - TN
    
    if pos_class == 1:
        return TP, FP, TN, FN
    else:
        return TN, FN, TP, FP