Doctorado-ML/stree
1
1
Fork 0
mirror of https://github.com/Doctorado-ML/STree.git synced 2025-08-15 15:36:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix compute number of nodes

Browse Source
This commit is contained in:
Ricardo Montañana Gómez
2021-04-13 22:31:05 +02:00
parent 783d105099
commit b55f59a3ec
2 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
stree/Strees.py
View File

@@ -841,10 +841,9 @@ class Stree(BaseEstimator, ClassifierMixin):
nodes = 0
leaves = 0
for node in self:
nodes += 1
if node.is_leaf():
leaves += 1
else:
nodes += 1
return nodes, leaves
def __iter__(self) -> Siterator:

6
stree/tests/Stree_test.py
View File

@@ -465,13 +465,13 @@ class Stree_test(unittest.TestCase):
clf = Stree(random_state=self._random_state)
clf.fit(X, y)
nodes, leaves = clf.nodes_leaves()
self.assertEqual(12, nodes)
self.assertEqual(25, nodes)
self.assertEquals(13, leaves)
X, y = load_wine(return_X_y=True)
clf = Stree(random_state=self._random_state)
clf.fit(X, y)
nodes, leaves = clf.nodes_leaves()
self.assertEqual(4, nodes)
self.assertEqual(9, nodes)
self.assertEquals(5, leaves)
def test_nodes_leaves_artificial(self):
@@ -489,5 +489,5 @@ class Stree_test(unittest.TestCase):
clf = Stree(random_state=self._random_state)
clf.tree_ = n1
nodes, leaves = clf.nodes_leaves()
self.assertEqual(4, nodes)
self.assertEqual(6, nodes)
self.assertEqual(2, leaves)