ddbourgin / numpy-ml

Machine learning, in numpy
https://numpy-ml.readthedocs.io/
GNU General Public License v3.0
15.35k stars 3.72k forks source link

Possibly bug in decision tree #39

Closed jcl2018 closed 5 years ago

jcl2018 commented 5 years ago

I was trying to run decision tree model using the very simple Iris dataset and I found a possible bug:

In trees/dt.py, line 160: image

The problem is when the array "levels" only has one element, the threshold will be empty, and it will crush on line 164 when we call "gain.max()".

Please correct me if I'm wrong. Thank you for creating this useful repo!

ddbourgin commented 5 years ago

Hey @jcl2018 - thanks a lot for this report, and sorry for the slow response! You are, of course, absolutely right :). I'd be happy to accept a PR to fix this if you feel like writing one -- otherwise, I will try to fix this shortly. Thank you!