python_analysis/second_semestr/pract6-1.py - main - Nik_voz/rgpu_python_analysis

python_analysis/second_semestr/pract6-1.py

67 строк2 KB

Nick-voz

feat: add decision tree classifier and random forest regressor implementations

26 мар 2026, 06:50

26 мар 2026, 06:50c1feb41

import os

import pandas as pd
from constants import Columns
from dotenv.main import load_dotenv
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
from sklearn.metrics import r2_score
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import train_test_split

load_dotenv()
if (DATASET_PATH := os.environ.get("DATASET_PATH")) is None:
    raise RuntimeWarning("DATASET_PATH is not set in environment")

data = pd.read_csv(DATASET_PATH, sep=";")

data = data[data[Columns.AVERAGE_SALARY_FACT_AVG].notna()]
data = data[data[Columns.COUNT_GRADUATE] > 0]

data[Columns.COUNT_GRADUATE] = data[Columns.COUNT_GRADUATE].fillna(
    data[Columns.COUNT_GRADUATE].median()
)

data[Columns.PERCENT_COMBINE_WORK_AND_STUDY] = data[
    Columns.PERCENT_COMBINE_WORK_AND_STUDY
].fillna(0)

features = [
    Columns.GENDER,
    Columns.EDUCATION_LEVEL,
    Columns.YEAR,
    Columns.COUNT_GRADUATE,
    Columns.PERCENT_COMBINE_WORK_AND_STUDY,
]

y = data[Columns.AVERAGE_SALARY_FACT_AVG]
X = pd.get_dummies(
    data[features], columns=[Columns.GENDER, Columns.EDUCATION_LEVEL]
)

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Случайный лес выбран, так как он эффективно обрабатывает нелинейные зависимости
# и устойчив к выбросам в данных о зарплатах. Ожидаем высокую точность за счет ансамбля деревьев.
rf = RandomForestRegressor(random_state=42)

rf.fit(X_train, y_train)

y_pred = rf.predict(X_test)
mse = mean_squared_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

print(f"MSE: {mse}, R2: {r2}")

param_grid = {"n_estimators": [50, 100], "max_depth": [10, 20, None]}
grid_search = GridSearchCV(
    RandomForestRegressor(random_state=42), param_grid, cv=3
)
grid_search.fit(X_train, y_train)

final_mse = mean_squared_error(y_test, grid_search.predict(X_test))
final_r2 = r2_score(y_test, grid_search.predict(X_test))

print(f"MSE: {final_mse}, R2: {final_r2}")