?

案例代碼

一、代碼說明

本代碼針對員工離職預測問題，使用CART決策樹算法（基尼指數）實現分類，并包含特征重要性評估和樹結構可視化。數據為模擬的10個員工樣本，特征包括工作年限、月薪、是否加班、團隊氛圍評分，目標是預測是否離職。代碼邏輯清晰，適合數模小白入門。

二、完整代碼

# 導入必要的庫
import pandas as pd
import numpy as np
from sklearn.preprocessing import LabelEncoder  # 用于 categorical 特征轉數值
from sklearn.model_selection import train_test_split  # 劃分訓練集/測試集
from sklearn.tree import DecisionTreeClassifier, plot_tree  # 決策樹模型及可視化
from sklearn.metrics import accuracy_score  # 計算準確率
import matplotlib.pyplot as plt  # 繪圖庫def prepare_data():"""生成并預處理模擬數據（修正后的數據：使“是否加班”與“是否離職”不完全相關）返回：特征矩陣X，目標向量y，特征名稱feature_names，類別名稱class_names"""# 1. 模擬員工離職數據（修正：調整2個樣本的離職狀態，使特征與目標不完全相關）data = pd.DataFrame({'工作年限（年）': [1, 3, 2, 5, 1, 4, 2, 3, 1, 5],'月薪（元）': [5000, 8000, 6000, 10000, 5500, 9000, 6500, 7000, 4500, 11000],'是否加班': ['是', '否', '是', '否', '是', '否', '否', '是', '是', '否'],'團隊氛圍評分（1-5）': [3, 4, 2, 5, 3, 4, 3, 2, 1, 5],'是否離職': ['是', '否', '是', '是', '是', '否', '否', '否', '是', '否']  # 修正：索引3（不加班）離職，索引7（加班）不離職})# 2. 處理 categorical 特征（將文字轉為數值，決策樹需要數值輸入）le_overtime = LabelEncoder()  # 處理“是否加班”data['是否加班'] = le_overtime.fit_transform(data['是否加班'])  # 否→0，是→1le_leave = LabelEncoder()  # 處理“是否離職”（目標變量）data['是否離職'] = le_leave.fit_transform(data['是否離職'])  # 否→0，是→1# 3. 劃分特征（X）和目標（y）feature_names = ['工作年限（年）', '月薪（元）', '是否加班', '團隊氛圍評分（1-5）']X = data[feature_names]y = data['是否離職']# 4. 返回預處理后的數據return X, y, feature_names, le_leave.classes_  # classes_是類別名稱（['否', '是']）def train_and_evaluate(X, y, feature_names, max_depth=3):"""訓練決策樹模型并評估性能參數：X（特征）、y（目標）、feature_names（特征名稱）、max_depth（樹最大深度，防止過擬合）返回：訓練好的模型clf，測試集準確率accuracy"""# 1. 劃分訓練集（80%）和測試集（20%）X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42  # random_state固定隨機種子，結果可重復)# 2. 初始化決策樹分類器（CART算法，使用基尼指數）clf = DecisionTreeClassifier(criterion='gini',  # 分裂準則：基尼指數（衡量混亂程度）max_depth=max_depth,  # 預剪枝：限制樹的最大深度，防止過擬合random_state=42)# 3. 訓練模型clf.fit(X_train, y_train)# 4. 用測試集預測并評估準確率y_pred = clf.predict(X_test)accuracy = accuracy_score(y_test, y_pred)# 5. 輸出結果print(f"測試集準確率：{accuracy:.2f}")  # 打印準確率（保留兩位小數）return clf, accuracydef plot_feature_importance(clf, feature_names):"""可視化特征重要性（決策樹認為哪些特征對預測最有用）參數：clf（訓練好的模型）、feature_names（特征名稱）"""# 1. 獲取特征重要性（數值越大，特征越重要）importances = clf.feature_importances_# 2. 繪制柱狀圖plt.figure(figsize=(10, 6))  # 設置圖的大小plt.bar(feature_names,  # x軸：特征名稱importances,    # y軸：重要性得分color='skyblue' # 柱子顏色)plt.xlabel('特征名稱', fontsize=12)  # x軸標簽plt.ylabel('重要性得分', fontsize=12)  # y軸標簽plt.title('員工離職預測特征重要性', fontsize=14)  # 圖標題plt.xticks(rotation=45)  # 旋轉x軸標簽，避免重疊plt.tight_layout()  # 自動調整布局，防止標簽截斷plt.show()  # 顯示圖def plot_tree_structure(clf, feature_names, class_names):"""可視化決策樹結構（直觀看到決策過程）參數：clf（訓練好的模型）、feature_names（特征名稱）、class_names（類別名稱，如['否', '是']）"""# 1. 設置圖的大小plt.figure(figsize=(20, 12))# 2. 繪制決策樹plot_tree(clf,                      # 訓練好的模型feature_names=feature_names,  # 特征名稱（顯示在節點上）class_names=class_names,      # 類別名稱（顯示在葉節點上，如['否', '是']）filled=True,               # 用顏色填充節點（顏色越深，節點越純）rounded=True,              # 節點邊框用圓角fontsize=10,               # 字體大小proportion=True            # 顯示樣本比例（如“3/5”表示該節點有3個正例，5個總樣本）)# 3. 設置圖標題plt.title('員工離職預測決策樹結構', fontsize=14)# 4. 顯示圖plt.show()# 主程序（入口函數）
if __name__ == "__main__":# 1. 準備數據X, y, feature_names, class_names = prepare_data()print("數據預處理完成！")# 2. 訓練模型并評估print("\n=== 模型訓練與評估 ===")clf, accuracy = train_and_evaluate(X, y, feature_names, max_depth=3)# 3. 特征重要性分析print("\n=== 特征重要性評估 ===")plot_feature_importance(clf, feature_names)# 4. 決策樹結構可視化print("\n=== 決策樹結構可視化 ===")plot_tree_structure(clf, feature_names, class_names)

三、代碼使用說明

1.環境準備

需要安裝以下Python庫（用pip命令安裝）：

pPip?install?pandas?numpy?scikit-learn?matplotlib

2.運行代碼

將代碼保存為employee_turnover_prediction.py，在終端或IDE中運行：

python?employee_turnover_prediction.py

3.輸出解釋

數據預處理完成：提示數據生成和轉換成功（將“是否加班”“是否離職”轉為0/1）。

模型訓練與評估：輸出測試集準確率（如測試集準確率：0.50或1.00，取決于測試集樣本分布）。

特征重要性評估：彈出柱狀圖，顯示每個特征對預測的影響（數值越大，特征越重要，修正后“是否加班”仍為關鍵特征，但需結合其他特征判斷）。

決策樹結構可視化：彈出決策樹圖，解讀方式如下：

根節點：第一個決策（如“是否加班≤0.5”，即“否”），分裂后左邊是“否”的樣本（部分離職），右邊是“是”的樣本（部分離職）。

內部節點：中間決策（如“團隊氛圍評分（1-5）≤2.5”），分裂后左邊是評分≤2的樣本（高離職率），右邊是評分≥3的樣本（低離職率）。

葉節點：最終結論（如“否”表示不離職，“是”表示離職），節點中的數字表示樣本比例（如“1/3”表示該節點有1個正例，3個總樣本）。

4.調參建議

max_depth：樹的最大深度（默認3），如果過擬合（訓練準確率100%，測試準確率低），可以減小max_depth；如果欠擬合（準確率低），可以增大max_depth。

criterion：分裂準則（默認gini），可以嘗試entropy（信息增益），兩者結果通常一致。

四、關鍵結論（修正后）

特征重要性：“是否加班”仍是預測離職的最關鍵特征，但“團隊氛圍評分”和“月薪”也能輔助區分（如加班但團隊氛圍好的員工可能不離職）。

決策邏輯（以修正后的數據為例）：

首先判斷“是否加班”：不加班的員工中，部分因其他原因離職（如索引3的高月薪員工）。

對于加班的員工，再判斷“團隊氛圍評分”：評分≤2的員工離職率高（如索引2、8的低評分員工），評分≥3的員工需進一步判斷。

對于加班且團隊氛圍評分≥3的員工，最后判斷“月薪”：月薪≤5750元的員工離職（如索引0、4的低薪員工），月薪>5750元的員工可能不離職（如索引7的高薪員工）。

通過本代碼，你可以快速掌握決策樹的實現流程，并直觀看到決策樹如何“思考”問題，非常適合數模比賽中的分類問題！