gavin_resch
gavin_resch

Reputation: 1

Remove empty outputs from scikit-learn KDtree.query_radius() and get unique values

I'm trying to flatten the output of a scikit-learn KDtree. I'd also like to call np.unique() on the flattened array.

The call I'm making is:

kd_tree = KDTree(X)
idxs = kd_tree.query_radius(Y, r=0.2)

The dimensions of X, Y are (49,2) (they are 2D coordinates).

The output I am getting is:

[array([], dtype=int64) array([7]) array([ 8, 21]) array([ 9, 22, 35])
 array([10, 23, 36]) array([11, 24, 37]) array([12, 25, 38]) array([0])
 array([ 1, 14]) array([ 2, 15, 28]) array([ 3, 16, 29, 42])
 array([ 4, 17, 30, 43]) array([ 5, 18, 31, 44]) array([ 6, 19, 32, 45])
 array([7]) array([ 8, 21]) array([ 9, 22, 35]) array([10, 23, 36])
 array([11, 24, 37]) array([12, 25, 38]) array([13, 26, 39])
 array([ 1, 14]) array([ 2, 15, 28]) array([ 3, 16, 29, 42])
 array([ 4, 17, 30, 43]) array([ 5, 18, 31, 44]) array([ 6, 19, 32, 45])
 array([20, 33, 46]) array([ 8, 21]) array([ 9, 22, 35])
 array([10, 23, 36]) array([11, 24, 37]) array([12, 25, 38])
 array([13, 26, 39]) array([27, 40]) array([ 2, 15, 28])
 array([ 3, 16, 29, 42]) array([ 4, 17, 30, 43]) array([ 5, 18, 31, 44])
 array([ 6, 19, 32, 45]) array([20, 33, 46]) array([34, 47])
 array([ 9, 22, 35]) array([10, 23, 36]) array([11, 24, 37])
 array([12, 25, 38]) array([13, 26, 39]) array([27, 40]) array([41])]

I would like to the output to be (without having to loop over the array):

[0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47]

Upvotes: 0

Views: 24

Answers (1)

James
James

Reputation: 36721

If the output you have shown is idxs, you can use hstack and unique from numpy.

np.unique(np.hstack(idxs))

Upvotes: 0

Related Questions