带赋值操作的列表理解

Question

我的代码没有错误。但我有点想知道。有没有办法使用列表理解来编写我的代码？我想删除外循环（for i in range(labels)）并用列表理解编写它。实际上，我面临着如何在内部和外部循环之间进行分配（current_class_p = p[y == i]）的问题。

例如，y = np.array([0, 1, 1, 0, 1])、p =np.array([1, 0, 1, 0, 1]) 和此 [[1, 1],[1, 2]] 的混淆矩阵。

def confusion_matrix_version2(y, p):
  labels = len(np.unique(y)) 
  result = np.zeros((labels, labels), dtype=int)

  for i in range(labels):
    current_class_p = p[y == i]
    result[i] = [len(current_class_p[current_class_p == j]) for j in range(labels)]
  return result

Answer 1

使用sklearn 的confusion_matrix 生成您想要的输出：

>>> from sklearn.metrics import confusion_matrix
>>> y_true = [2, 0, 2, 2, 0, 1]
>>> y_pred = [0, 0, 2, 2, 0, 2]
>>> cm = confusion_matrix(y_true, y_pred)
array([[2, 0, 0],
       [0, 0, 1],
       [1, 0, 2]])

其中行对应于预测值，列对应于实际值。如果你想把它作为一个嵌套列表，你可以使用cm.tolist()

>>> cm.tolist()
[[2, 0, 0], [0, 0, 1], [1, 0, 2]]

编辑：根据 juanpa.arrivillaga 的建议，将列表推导中的列表输出更新为数组的 tolist 函数。