Manage corner case with empty column of labels in predict

scikit-learn · Aug 13, 2014 · 043ed2e · 043ed2e
1 parent bb1c0ae
commit 043ed2e
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/sklearn/neighbors/classification.py b/sklearn/neighbors/classification.py
@@ -179,10 +179,12 @@ def predict(self, X):
                 # Find the neigh_ind in _y.data using _y.indices as a guide
                 data_index = np.searchsorted(_y_indices_k, neigh_ind)
                 data_index[data_index == _y_data_k.shape[0]] = 0
-                neigh_lbls_k = _y_data_k[data_index]
-
-                # Replace incorrect nonzero elements with correct zeros
-                neigh_lbls_k[_y_indices_k[data_index] != neigh_ind] = 0
+                if _y_data_k.size == 0:
+                    neigh_lbls_k = np.zeros(shape=data_index.shape)
+                else:
+                    neigh_lbls_k = _y_data_k[data_index]
+                    # Replace incorrect nonzero elements with correct zeros
+                    neigh_lbls_k[_y_indices_k[data_index] != neigh_ind] = 0
 
                 if weights is None:
                     mode = csr_row_mode(sp.csr_matrix(neigh_lbls_k))