From c8b09a4e363c38b5338896623f22ac9df29d3d66 Mon Sep 17 00:00:00 2001
From: Mikhail Martin <JohnDoe110@yandex.ru>
Date: Sat, 27 Sep 2025 00:00:03 +0300
Subject: [PATCH] optimized

---
 smarttree/_builder.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/smarttree/_builder.py b/smarttree/_builder.py
index c23c222..346acd3 100644
--- a/smarttree/_builder.py
+++ b/smarttree/_builder.py
@@ -100,9 +100,15 @@ def build(self, tree: Tree) -> None:
             tree.leaf_counter -= 1
 
     def distribution(self, mask: pd.Series) -> NDArray[np.integer]:
-        return np.array([
-            (mask & (self.y == class_name)).sum() for class_name in self.class_names
-        ])
+
+        mask_arr = mask.to_numpy()
+        y_arr = self.y.to_numpy()
+
+        result = np.zeros(len(self.class_names), dtype=np.int32)
+        for i, class_name in enumerate(self.class_names):
+            result[i] = np.sum(mask_arr & (y_arr == class_name))
+
+        return result
 
     def gini_index(self, mask: pd.Series) -> float:
         r"""