Mean median and other function effect all data and now only column, removing useless columns

src/back/managing_missing_values.py

@@ -19,34 +19,28 @@ def drop_high_null_percentage(data, threshold=0.5):
     return data
 
 
+def replace_with_mean(data):
+    return data.apply(lambda col: col.fillna(col.mean()) if col.dtype.kind in 'biufc' else col)
 
-def replace_with_mean(data, column):
-    data[column] = data[column].fillna(data[column].mean())
-    return data
-
-def replace_with_median(data, column):
-    data[column] = data[column].fillna(data[column].median())
-    return data
+def replace_with_median(data):
+    return data.apply(lambda col: col.fillna(col.median()) if col.dtype.kind in 'biufc' else col)
 
-def replace_with_mode(data, column):
-    mode_value = data[column].mode()
-    if not mode_value.empty:
-        data[column] = data[column].fillna(mode_value[0])
-    return data
+def replace_with_mode(data):
+    return data.apply(lambda col: col.fillna(col.mode()[0]) if col.mode().size > 0 else col)
 
-def impute_with_knn(data, column, n_neighbors=5):
+def impute_with_knn(data, n_neighbors=5):
     imputer = KNNImputer(n_neighbors=n_neighbors)
-    data[[column]] = imputer.fit_transform(data[[column]])
-    return data
+    return pd.DataFrame(imputer.fit_transform(data), columns=data.columns)
 
-def impute_with_regression(data, column):
-    if data[column].isnull().sum() > 0:
-        train_data = data[data[column].notna()]
-        test_data = data[data[column].isna()]
-        if not train_data.empty and not test_data.empty:
-            regressor = LinearRegression()
-            regressor.fit(train_data.drop(column, axis=1), train_data[column])
-            data.loc[data[column].isna(), column] = regressor.predict(test_data.drop(column, axis=1))
+def impute_with_regression(data):
+    for column in data.columns:
+        if data[column].isnull().sum() > 0:
+            train_data = data[data[column].notna()]
+            test_data = data[data[column].isna()]
+            if not train_data.empty and not test_data.empty:
+                regressor = LinearRegression()
+                regressor.fit(train_data.drop(column, axis=1), train_data[column])
+                data.loc[data[column].isna(), column] = regressor.predict(test_data.drop(column, axis=1))
     return data
 
 
@@ -56,28 +50,20 @@ def impute_with_regression(data, column):
     - method: Method to handle missing values ('drop', 'mean', 'median', 'mode', 'knn', 'regression')
     - n_neighbors: Number of neighbors to use for KNN imputation (only used if method='knn')
 """
-def handle_missing_values(data, method, column, n_neighbors=5):
+def handle_missing_values(data, method, n_neighbors=5):
 
     data = drop_high_null_percentage(data)
-    data = convert_categorical_to_numeric(data)
-        
+    data = convert_categorical_to_numeric(data)    
     if method == 'mean':
-        return replace_with_mean(data, column)
+        return replace_with_mean(data)
     elif method == 'median':
-        return replace_with_median(data, column)
+        return replace_with_median(data)
     elif method == 'mode':
-        return replace_with_mode(data, column)
+        return replace_with_mode(data)
     elif method == 'knn':
-        return impute_with_knn(data, column, n_neighbors=n_neighbors)
+        return impute_with_knn(data, n_neighbors=n_neighbors)
     elif method == 'regression':
-        return impute_with_regression(data, column)
-    elif method == 'drop_high_null':
-        return drop_high_null_percentage(data)
+        return impute_with_regression(data)
     else:
         raise ValueError("Unknown method")
 
-
-
-data = l.return_csv('./data.csv')
-cleaned_data = handle_missing_values(data, method='mode', column='Route Type')
-print(cleaned_data)