如何使用NB算法对指定人口数据集的built列分类并生成三类混淆矩阵

• 第一步：加载所需依赖包

# 未安装的话先执行 install.packages(c("e1071", "caret"))
library(e1071) # 内置朴素贝叶斯算法实现
library(caret) # 用于数据集拆分、混淆矩阵计算

• 第二步：拆分训练集与测试集

# 提取特征（你已经处理好的3-6列）和目标标签
features <- data[, c("height", "weight", "educated", "acceptchanges")]
labels <- data$built

# 固定随机种子保证结果可复现，按7:3比例拆分训练测试集
set.seed(123)
train_idx <- createDataPartition(labels, p = 0.7, list = FALSE)
train_x <- features[train_idx, ]
train_y <- labels[train_idx]
test_x <- features[-train_idx, ]
test_y <- labels[-train_idx]

• 第三步：训练朴素贝叶斯模型

# 可添加laplace = 1参数启用拉普拉斯平滑，避免零概率问题
nb_model <- naiveBayes(x = train_x, y = train_y, laplace = 1)

• 第四步：预测测试集结果

pred_y <- predict(nb_model, newdata = test_x)

• 第五步：输出三分类混淆矩阵

confusion_matrix <- confusionMatrix(pred_y, test_y)
# 直接打印即可查看三个类别的混淆矩阵及配套指标
print(confusion_matrix)

如果习惯用Python实现，可参考以下代码：

• 依赖导入

import pandas as pd
from sklearn.naive_bayes import GaussianNB
from sklearn.model_selection import train_test_split
from sklearn.metrics import confusion_matrix

• 拆分数据集与训练预测

# 提取特征和标签
X = data[["height", "weight", "educated", "acceptchanges"]]
y = data["built"]

# 分层拆分保证各类别样本占比一致
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, random_state=123, stratify=y
)

# 训练高斯朴素贝叶斯（适配连续型特征如身高、体重）
nb_model = GaussianNB()
nb_model.fit(X_train, y_train)
y_pred = nb_model.predict(X_test)

# 输出指定类别顺序的混淆矩阵
cm = confusion_matrix(y_test, y_pred, labels=["hillyreg", "plainsreg", "mountreg"])
print(cm)

小提示：如果你的布尔型特征educated、acceptchanges是离散分类特征，可将高斯朴素贝叶斯替换为多项式朴素贝叶斯MultinomialNB，效果会更好。

内容的提问来源于stack exchange，提问作者vicky