From 31d79a77fa82fb1bc0869cbf3eeec5814a5d8b5c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ricardo=20Monta=C3=B1ana?= <rmontanana@gmail.com>
Date: Mon, 13 Feb 2023 17:34:50 +0100
Subject: [PATCH] Add get_states_feature method

---
 src/fimdlp/mdlp.py | 26 +++++++++++++++++---------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/fimdlp/mdlp.py b/src/fimdlp/mdlp.py
index ac0dd05..808e75a 100644
--- a/src/fimdlp/mdlp.py
+++ b/src/fimdlp/mdlp.py
@@ -119,15 +119,6 @@ class FImdlp(TransformerMixin, BaseEstimator):
         else:
             result[:, feature] = X
 
-    def range_features(self):
-        res = []
-        for i in range(self.n_features_in_):
-            if i in self.features_:
-                res.append(list(range(len(self.cut_points_[i]))))
-            else:
-                res.append([])
-        return res
-
     def transform(self, X):
         """Discretize X values.
         Parameters
@@ -186,6 +177,23 @@ class FImdlp(TransformerMixin, BaseEstimator):
             result.append(self.cut_points_[feature])
         return result
 
+    def get_states_feature(self, feature):
+        """Return the states a feature can take
+
+        Parameters
+        ----------
+        feature : int
+            feature to get the states
+
+        Returns
+        -------
+        list
+            states of the feature
+        """
+        if feature in self.features_:
+            return list(range(len(self.cut_points_[feature]) + 1))
+        return None
+
     def join_fit(self, features, target, data):
         """Join the selected features with the labels and fit the discretizer
         of the target variable