class TestKDTree2d(unittest.TestCase):
def setUp(self):
points = [(2, 3), (5, 4), (9, 6), (4, 7), (8, 1), (7, 2)]
self.tree = KDTree(points)
def test_traversal(self):
logger.info("PreOrder:")
self.tree.traversal(lambda n: logger.info(n), "preorder")
logger.info("InOrder:")
self.tree.traversal(lambda n: logger.info(n), "inorder")
logger.info("PostOrder:")
self.tree.traversal(lambda n: logger.info(n), "postorder")
def test_closest(self):
dist, node, count, nodes = self.tree.closest((2.3, 3.3))
logger.info("dist = {0}, node = {1}, visit_count = {2}".format(
dist, node, count))
logger.info("visited nodes = {0}".format(nodes))
self.assertTrue(np.all(node.point == (2, 3)),
"actual is {0}".format(node.point))
dist, node, count, nodes = self.tree.closest((2.0, 4.5))
logger.info("dist = {0}, node = {1}, visit_count = {2}".format(
dist, node, count))
logger.info("visited nodes = {0}".format(nodes))
self.assertTrue(np.all(node.point == (2, 3)),
"actual is {0}".format(node.point))
dist, node, count, nodes = self.tree.closest((6.9, 4.2))
logger.info("dist = {0}, node = {1}, visit_count = {2}".format(
dist, node, count))
logger.info("visited nodes = {0}".format(nodes))
self.assertTrue(np.all(node.point == (5, 4)))
def test_visualization(self):
self.tree.traversal(draw_point, 'preorder')
show_closest(self.tree, (2.3, 3.3), 'r')
show_closest(self.tree, (2.0, 4.5), 'g')
show_closest(self.tree, (6.8, 4.2), 'm')
plt.show()
class TestKDTree1d(unittest.TestCase):
def setUp(self):
np.random.seed(0)
points = np.random.randint(0, 50, 10)
self.tree = KDTree(np.transpose([points]))
def test_traversal(self):
logger.info("PreOrder:")
self.tree.traversal(lambda n: logger.info(n), "preorder")
logger.info("InOrder:")
self.tree.traversal(lambda n: logger.info(n), "inorder")
logger.info("PostOrder:")
self.tree.traversal(lambda n: logger.info(n), "postorder")
def test_closest(self):
dist, node, count, nodes = self.tree.closest([15])
logger.info("dist = {0}, node = {1}, visit_count = {2}".format(
dist, node, count))
logger.info("visited nodes = {0}".format(nodes))
