stefankoegl / kdtree

A Python implementation of a kd-tree
ISC License
365 stars 118 forks source link

Can't remove custom item with payload #54

Closed joric closed 1 year ago

joric commented 1 year ago

What am I missing here? Looks like remove doesn't work properly for custom items. Using tree = tree.remove( [2, 3, 4], node = r ) for identity check doesn't work either.

import kdtree

class Item(object):
    def __init__(self, x, y, z, data):
        self.coords = (x, y, z)
        self.data = data

    def __len__(self):
        return len(self.coords)

    def __getitem__(self, i):
        return self.coords[i]

    def __repr__(self):
        return 'Item({}, {}, {})'.format(self.coords[0], self.coords[1], self.coords[2], self.data)

tree = kdtree.create(dimensions=3)

tree.add ( Item(2, 3, 4, 'First') )

tree.add ( Item(3, 4, 5, 'Second') )

r = tree.search_nn( [1, 2, 3] )

print(r) # (<KDNode - Item(2, 3, 4)>, 3.0)

tree = tree.remove( [2, 3, 4] )

r = tree.search_nn( [1, 2, 3] )

print(r) # (<KDNode - Item(2, 3, 4)>, 3.0)
joric commented 1 year ago

Ok got it after studying sources. We need to define eq, and pass identity node, e.g. like this:

    def __eq__(self, other):
        return all(x==y for x,y in zip(self.coords,other.coords))
...
tree = tree.remove( r[0].data, node=r[0] )