Mean Squared Error calculation instead of Gini-CodePudding

I have an information criterion with help of Gini. How can I change my code for Mean Squared Error instead of Gini? Python, the random forest task.

def gini(labels):
    #  to count the number of objects for different classes
    classes = {}
    for label in labels:
        if label not in classes:
            classes[label] = 0
        classes[label]  = 1
    
    #  to calculate Gini criterion
    impurity = 1
    for label in classes:
        p = classes[label] / len(labels)
        impurity -= p ** 2
    return impurity

CodePudding user response：

Here is how you could modify the GINI function to calculate the MSE :

def mse(labels):
    # Calculate the predicted value for each class
    classes = {}
    for label in labels:
        if label not in classes:
            classes[label] = []
        classes[label].append(label)
    
    predicted_values = {}
    for label in classes:
        predicted_values[label] = sum(classes[label]) / len(classes[label])
    
    # Calculate the MSE
    squared_errors = []
    for label in labels:
        squared_error = (label - predicted_values[label]) ** 2
        squared_errors.append(squared_error)
    if len(squared_errors) > 0:
        mse = sum(squared_errors) / len(squared_errors)
    else:
        mse = 0
    return mse