solve smotenc

Author: glemaitre

imblearn/over_sampling/_smote.py

@@ -900,11 +900,12 @@ def __init__(
         )
         self.categorical_features = categorical_features
 
-    @staticmethod
-    def _check_X_y(X, y):
+    def _check_X_y(self, X, y):
         """Overwrite the checking to let pass some string for categorical
         features.
         """
+        # store the columns name to reconstruct a dataframe
+        self._columns = X.columns if hasattr(X, "loc") else None
         y, binarize_y = check_target_type(y, indicate_one_vs_all=True)
         X, y = check_X_y(X, y, accept_sparse=["csr", "csc"], dtype=None)
         return X, y, binarize_y

imblearn/over_sampling/tests/test_smote_nc.py

@@ -13,6 +13,7 @@
 
 from sklearn.datasets import make_classification
 from sklearn.utils._testing import assert_allclose
+from sklearn.utils._testing import assert_array_equal
 
 from imblearn.over_sampling import SMOTENC
 
@@ -184,7 +185,7 @@ def test_smotenc_pandas():
     smote = SMOTENC(categorical_features=categorical_features, random_state=0)
     X_res_pd, y_res_pd = smote.fit_resample(X_pd, y)
     X_res, y_res = smote.fit_resample(X, y)
-    assert X_res_pd.tolist() == X_res.tolist()
+    assert_array_equal(X_res_pd.to_numpy(), X_res)
     assert_allclose(y_res_pd, y_res)