├── requirements-dev.txt ├── setup.cfg ├── AUTHORS ├── MANIFEST.in ├── .gitignore ├── .travis.yml ├── LICENSE ├── setup.py ├── readme.md ├── doc ├── index.rst ├── Makefile └── conf.py ├── bounded_priority_queue.py ├── test.py └── kdtree.py /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | wheel 2 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [bdist_wheel] 2 | universal=1 3 | -------------------------------------------------------------------------------- /AUTHORS: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Zverik/kdtree/master/AUTHORS -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include AUTHORS 2 | include LICENSE 3 | include readme.md 4 | include test.py 5 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | .coverage 3 | build 4 | dist 5 | MANIFEST 6 | doc/_build 7 | venv* 8 | kdtree.egg-info 9 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | - "2.6" 4 | - "2.7" 5 | - "3.3" 6 | - "3.4" 7 | - "3.5" 8 | - "nightly" 9 | - "pypy" 10 | - "pypy3" 11 | 12 | script: python test.py 13 | 14 | sudo: false 15 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Copyright (c) Stefan Kögl 3 | 4 | Permission to use, copy, modify, and/or distribute this software for any 5 | purpose with or without fee is hereby granted, provided that the above 6 | copyright notice and this permission notice appear in all copies. 7 | 8 | THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH 9 | REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND 10 | FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, 11 | INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM 12 | LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR 13 | OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR 14 | PERFORMANCE OF THIS SOFTWARE. 15 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from setuptools import setup 4 | import re 5 | import os.path 6 | 7 | dirname = os.path.dirname(os.path.abspath(__file__)) 8 | filename = os.path.join(dirname, 'kdtree.py') 9 | src = open(filename).read() 10 | metadata = dict(re.findall("__([a-z]+)__ = u?'([^']+)'", src)) 11 | docstrings = re.findall('"""([^"]+)"""', src) 12 | 13 | PACKAGE = 'kdtree' 14 | 15 | MODULES = ( 16 | PACKAGE, 17 | 'bounded_priority_queue', 18 | ) 19 | 20 | AUTHOR_EMAIL = metadata['author'] 21 | VERSION = metadata['version'] 22 | WEBSITE = metadata['website'] 23 | LICENSE = metadata['license'] 24 | DESCRIPTION = docstrings[0].strip() 25 | if '\n\n' in DESCRIPTION: 26 | DESCRIPTION, LONG_DESCRIPTION = DESCRIPTION.split('\n\n', 1) 27 | else: 28 | LONG_DESCRIPTION = None 29 | 30 | # Extract name and e-mail ("Firstname Lastname ") 31 | AUTHOR, EMAIL = re.match(r'(.*) <(.*)>', AUTHOR_EMAIL).groups() 32 | 33 | CLASSIFIERS = [ 34 | 'Development Status :: 5 - Production/Stable', 35 | 'Intended Audience :: Developers', 36 | 'License :: OSI Approved :: ISC License (ISCL)', 37 | 'Operating System :: OS Independent', 38 | 'Programming Language :: Python', 39 | 'Programming Language :: Python :: 2', 40 | 'Programming Language :: Python :: 2.6', 41 | 'Programming Language :: Python :: 2.7', 42 | 'Programming Language :: Python :: 3', 43 | 'Programming Language :: Python :: 3.3', 44 | 'Programming Language :: Python :: 3.4', 45 | 'Programming Language :: Python :: 3.5', 46 | 'Programming Language :: Python :: Implementation :: CPython', 47 | 'Programming Language :: Python :: Implementation :: PyPy', 48 | 'Topic :: Software Development :: Libraries', 49 | 'Topic :: Utilities', 50 | ] 51 | 52 | 53 | setup(name=PACKAGE, 54 | version=VERSION, 55 | description=DESCRIPTION, 56 | long_description=LONG_DESCRIPTION, 57 | author=AUTHOR, 58 | author_email=EMAIL, 59 | license=LICENSE, 60 | url=WEBSITE, 61 | py_modules=MODULES, 62 | download_url='http://pypi.python.org/packages/source/' + \ 63 | PACKAGE[0] + '/' + PACKAGE + '/' + \ 64 | PACKAGE + '-' + VERSION + '.tar.gz', 65 | classifiers=CLASSIFIERS, 66 | ) 67 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | A simple kd-tree in Python [![Build Status](https://travis-ci.org/stefankoegl/kdtree.png?branch=master)](https://travis-ci.org/stefankoegl/kdtree) 2 | ========================== 3 | 4 | The kdtree package can construct, modify and search 5 | [kd-trees](http://en.wikipedia.org/wiki/Kd-tree). 6 | 7 | * Website: https://github.com/stefankoegl/kdtree 8 | * Repository: https://github.com/stefankoegl/kdtree.git 9 | * Documentation: https://python-kdtree.readthedocs.org/ 10 | * PyPI: https://pypi.python.org/pypi/kdtree 11 | * Travis-CI: https://travis-ci.org/stefankoegl/kdtree 12 | * Coveralls: https://coveralls.io/r/stefankoegl/kdtree 13 | 14 | 15 | Usage 16 | ----- 17 | 18 | >>> import kdtree 19 | 20 | # Create an empty tree by specifying the number of 21 | # dimensions its points will have 22 | >>> emptyTree = kdtree.create(dimensions=3) 23 | 24 | # A kd-tree can contain different kinds of points, for example tuples 25 | >>> point1 = (2, 3, 4) 26 | 27 | # Lists can also be used as points 28 | >>> point2 = [4, 5, 6] 29 | 30 | # Other objects that support indexing can be used, too 31 | >>> import collections 32 | >>> Point = collections.namedtuple('Point', 'x y z') 33 | >>> point3 = Point(5, 3, 2) 34 | 35 | # A tree is created from a list of points 36 | >>> tree = kdtree.create([point1, point2, point3]) 37 | 38 | # Each (sub)tree is represented by its root node 39 | >>> tree 40 | 41 | 42 | # Adds a tuple to the tree 43 | >>> tree.add( (5, 4, 3) ) 44 | 45 | # Removes the previously added point and returns the new root 46 | >>> tree = tree.remove( (5, 4, 3) ) 47 | 48 | # Retrieving the Tree in inorder 49 | >>> list(tree.inorder()) 50 | [, , ] 51 | 52 | # Retrieving the Tree in level order 53 | >>> list(kdtree.level_order(tree)) 54 | [, , ] 55 | 56 | # Find the nearest node to the location (1, 2, 3) 57 | >>> tree.search_nn( (1, 2, 3) ) 58 | 59 | 60 | # Add a point to make the tree more interesting 61 | >>> tree.add( (10, 2, 1) ) 62 | 63 | # Visualize the Tree 64 | >>> kdtree.visualize(tree) 65 | 66 | 67 | [4, 5, 6] 68 | 69 | (2, 3, 4) Point(x=5, y=3, z=2) 70 | 71 | (10, 2, 1) 72 | 73 | # Take the right subtree of the root 74 | >>> subtree = tree.right 75 | 76 | # and detatch it 77 | >>> tree.right = None 78 | >>> kdtree.visualize(tree) 79 | 80 | [4, 5, 6] 81 | 82 | (2, 3, 4) 83 | 84 | >>> kdtree.visualize(subtree) 85 | 86 | Point(x=5, y=3, z=2) 87 | 88 | (10, 2, 1) 89 | 90 | # and re-attach it 91 | >>> tree.right = subtree 92 | >>> kdtree.visualize(tree) 93 | 94 | [4, 5, 6] 95 | 96 | (2, 3, 4) Point(x=5, y=3, z=2) 97 | 98 | (10, 2, 1) 99 | 100 | # Add a node to make the tree unbalanced 101 | >>> tree.is_balanced 102 | True 103 | >>> tree.add( (6, 1, 5) ) 104 | >>> tree.is_balanced 105 | False 106 | >>> kdtree.visualize(tree) 107 | 108 | [4, 5, 6] 109 | 110 | (2, 3, 4) Point(x=5, y=3, z=2) 111 | (10, 2, 1) 112 | (6, 1, 5) 113 | # rebalance the tree 114 | >>> tree = tree.rebalance() 115 | >>> tree.is_balanced 116 | True 117 | >>> kdtree.visualize(tree) 118 | 119 | Point(x=5, y=3, z=2) 120 | 121 | [4, 5, 6] (6, 1, 5) 122 | 123 | (2, 3, 4) 124 | -------------------------------------------------------------------------------- /doc/index.rst: -------------------------------------------------------------------------------- 1 | .. python-kdtree documentation master file, created by 2 | sphinx-quickstart on Sat Mar 22 15:55:01 2014. 3 | You can adapt this file completely to your liking, but it should at least 4 | contain the root `toctree` directive. 5 | 6 | python-kdtree 7 | ============= 8 | 9 | The ``kdtree`` package can construct, modify and search `kd-trees 10 | `_. 11 | 12 | Example Usage 13 | ------------- 14 | 15 | .. code-block:: python 16 | 17 | >>> import kdtree 18 | 19 | # Create an empty tree by specifying the number of 20 | # dimensions its points will have 21 | >>> emptyTree = kdtree.create(dimensions=3) 22 | 23 | # A kd-tree can contain different kinds of points, for example tuples 24 | >>> point1 = (2, 3, 4) 25 | 26 | # Lists can also be used as points 27 | >>> point2 = [4, 5, 6] 28 | 29 | # Other objects that support indexing can be used, too 30 | >>> import collections 31 | >>> Point = collections.namedtuple('Point', 'x y z') 32 | >>> point3 = Point(5, 3, 2) 33 | 34 | # A tree is created from a list of points 35 | >>> tree = kdtree.create([point1, point2, point3]) 36 | 37 | # Each (sub)tree is represented by its root node 38 | >>> tree 39 | 40 | 41 | # Adds a tuple to the tree 42 | >>> tree.add( (5, 4, 3) ) 43 | 44 | # Removes the previously added point and returns the new root 45 | >>> tree = tree.remove( (5, 4, 3) ) 46 | 47 | # Retrieving the Tree in inorder 48 | >>> list(tree.inorder()) 49 | [, , ] 50 | 51 | # Retrieving the Tree in level order 52 | >>> list(kdtree.level_order(tree)) 53 | [, , ] 54 | 55 | # Find the nearest node to the location (1, 2, 3) 56 | >>> tree.search_nn( (1, 2, 3) ) 57 | 58 | 59 | # Add a point to make the tree more interesting 60 | >>> tree.add( (10, 2, 1) ) 61 | 62 | # Visualize the Tree 63 | >>> kdtree.visualize(tree) 64 | 65 | 66 | [4, 5, 6] 67 | 68 | (2, 3, 4) Point(x=5, y=3, z=2) 69 | 70 | (10, 2, 1) 71 | 72 | # Take the right subtree of the root 73 | >>> subtree = tree.right 74 | 75 | # and detatch it 76 | >>> tree.right = None 77 | >>> kdtree.visualize(tree) 78 | 79 | [4, 5, 6] 80 | 81 | (2, 3, 4) 82 | 83 | >>> kdtree.visualize(subtree) 84 | 85 | Point(x=5, y=3, z=2) 86 | 87 | (10, 2, 1) 88 | 89 | # and re-attach it 90 | >>> tree.right = subtree 91 | >>> kdtree.visualize(tree) 92 | 93 | [4, 5, 6] 94 | 95 | (2, 3, 4) Point(x=5, y=3, z=2) 96 | 97 | (10, 2, 1) 98 | 99 | # Add a node to make the tree unbalanced 100 | >>> tree.is_balanced 101 | True 102 | >>> tree.add( (6, 1, 5) ) 103 | >>> tree.is_balanced 104 | False 105 | >>> kdtree.visualize(tree) 106 | 107 | [4, 5, 6] 108 | 109 | (2, 3, 4) Point(x=5, y=3, z=2) 110 | (10, 2, 1) 111 | (6, 1, 5) 112 | # rebalance the tree 113 | >>> tree = tree.rebalance() 114 | >>> tree.is_balanced 115 | True 116 | >>> kdtree.visualize(tree) 117 | 118 | Point(x=5, y=3, z=2) 119 | 120 | [4, 5, 6] (6, 1, 5) 121 | 122 | (2, 3, 4) 123 | 124 | 125 | .. toctree:: 126 | :maxdepth: 2 127 | 128 | 129 | Indices and tables 130 | ================== 131 | 132 | * :ref:`genindex` 133 | * :ref:`modindex` 134 | * :ref:`search` 135 | 136 | -------------------------------------------------------------------------------- /bounded_priority_queue.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # 3 | # Based on the work of Pravin Paratey (April 15, 2011) 4 | # Joachim Hagege (June 18, 2014) 5 | # 6 | # Code released under BSD license 7 | # 8 | DISTANCE_INDEX = 1 9 | 10 | class BoundedPriorityQueue: 11 | """ This class illustrates a PriorityQueue and its associated functions """ 12 | 13 | def __init__(self, k): 14 | self.heap = [] 15 | self.k = k 16 | 17 | def items(self): 18 | return self.heap 19 | 20 | def parent(self, index): 21 | """ 22 | Parent will be at math.floor(index/2). 23 | """ 24 | return int(index / 2) 25 | 26 | def left_child(self, index): 27 | return 2 * index + 1 28 | 29 | def right_child(self, index): 30 | return 2 * index + 2 31 | 32 | def max_heapify(self, index): 33 | """ 34 | Responsible for maintaining the heap property of the heap. 35 | This function assumes that the subtree located at left and right 36 | child satisfies the max-heap property. But the tree at index 37 | (current node) does not. O(log n) 38 | """ 39 | left_index = self.left_child(index) 40 | right_index = self.right_child(index) 41 | 42 | largest = index 43 | if left_index < len(self.heap) and self._dist(left_index) > self._dist(index): 44 | largest = left_index 45 | if right_index < len(self.heap) and self._dist(right_index) > self._dist(largest): 46 | largest = right_index 47 | 48 | if largest != index: 49 | self.heap[index], self.heap[largest] = self.heap[largest], self.heap[index] 50 | self.max_heapify(largest) 51 | 52 | def _dist(self, index): 53 | """ Get the distance of the heap object at the given index """ 54 | return self.heap[index][DISTANCE_INDEX] 55 | 56 | def propagate_up(self, index): 57 | """ Compares index with parent and swaps node if larger O(log(n)) """ 58 | while index != 0 and self._dist(self.parent(index)) < self._dist(index): 59 | self.heap[index], self.heap[self.parent(index)] = self.heap[self.parent(index)], self.heap[index] 60 | index = self.parent(index) 61 | 62 | def add(self, obj): 63 | """ 64 | Add obj to the priority queue if it has a lower value than 65 | the maximum already in the queue. If the queue is full, 66 | then the object with the maximum value on the queue is 67 | replaced by this one. 68 | """ 69 | size = self.size() 70 | 71 | # Size == k, The priority queue is at capacity. 72 | if size == self.k: 73 | max_elem = self.max() 74 | 75 | # The new element has a lower distance than the biggest one. 76 | # Then we insert, otherwise, don't insert. 77 | if obj[DISTANCE_INDEX] < max_elem: 78 | self.extract_max() 79 | self.heap_append(obj) 80 | 81 | # if size == 0 or 0 < Size < k 82 | else: 83 | self.heap_append(obj) 84 | 85 | def heap_append(self, obj): 86 | """ Adds an element in the heap O(ln(n)) """ 87 | self.heap.append(obj) 88 | # Index value is 1 less than length: 89 | self.propagate_up(len(self.heap) - 1) 90 | 91 | def max(self): 92 | # The highest distance will always be at the index 0 (heap invariant) 93 | return self.heap[0][1] 94 | 95 | def size(self): 96 | return len(self.heap) 97 | 98 | def extract_max(self): 99 | """ 100 | Part of the Priority Queue, extracts the element on the top of the heap 101 | and then re-heapifies. O(log n). 102 | """ 103 | max = self.heap[0] 104 | data = self.heap.pop() 105 | if len(self.heap) > 0: 106 | self.heap[0] = data 107 | self.max_heapify(0) 108 | return max 109 | -------------------------------------------------------------------------------- /doc/Makefile: -------------------------------------------------------------------------------- 1 | # Makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line. 5 | SPHINXOPTS = 6 | SPHINXBUILD = sphinx-build 7 | PAPER = 8 | BUILDDIR = _build 9 | 10 | # User-friendly check for sphinx-build 11 | ifeq ($(shell which $(SPHINXBUILD) >/dev/null 2>&1; echo $$?), 1) 12 | $(error The '$(SPHINXBUILD)' command was not found. Make sure you have Sphinx installed, then set the SPHINXBUILD environment variable to point to the full path of the '$(SPHINXBUILD)' executable. Alternatively you can add the directory with the executable to your PATH. If you don't have Sphinx installed, grab it from http://sphinx-doc.org/) 13 | endif 14 | 15 | # Internal variables. 16 | PAPEROPT_a4 = -D latex_paper_size=a4 17 | PAPEROPT_letter = -D latex_paper_size=letter 18 | ALLSPHINXOPTS = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) . 19 | # the i18n builder cannot share the environment and doctrees with the others 20 | I18NSPHINXOPTS = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) . 21 | 22 | .PHONY: help clean html dirhtml singlehtml pickle json htmlhelp qthelp devhelp epub latex latexpdf text man changes linkcheck doctest gettext 23 | 24 | help: 25 | @echo "Please use \`make ' where is one of" 26 | @echo " html to make standalone HTML files" 27 | @echo " dirhtml to make HTML files named index.html in directories" 28 | @echo " singlehtml to make a single large HTML file" 29 | @echo " pickle to make pickle files" 30 | @echo " json to make JSON files" 31 | @echo " htmlhelp to make HTML files and a HTML help project" 32 | @echo " qthelp to make HTML files and a qthelp project" 33 | @echo " devhelp to make HTML files and a Devhelp project" 34 | @echo " epub to make an epub" 35 | @echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter" 36 | @echo " latexpdf to make LaTeX files and run them through pdflatex" 37 | @echo " latexpdfja to make LaTeX files and run them through platex/dvipdfmx" 38 | @echo " text to make text files" 39 | @echo " man to make manual pages" 40 | @echo " texinfo to make Texinfo files" 41 | @echo " info to make Texinfo files and run them through makeinfo" 42 | @echo " gettext to make PO message catalogs" 43 | @echo " changes to make an overview of all changed/added/deprecated items" 44 | @echo " xml to make Docutils-native XML files" 45 | @echo " pseudoxml to make pseudoxml-XML files for display purposes" 46 | @echo " linkcheck to check all external links for integrity" 47 | @echo " doctest to run all doctests embedded in the documentation (if enabled)" 48 | 49 | clean: 50 | rm -rf $(BUILDDIR)/* 51 | 52 | html: 53 | $(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html 54 | @echo 55 | @echo "Build finished. The HTML pages are in $(BUILDDIR)/html." 56 | 57 | dirhtml: 58 | $(SPHINXBUILD) -b dirhtml $(ALLSPHINXOPTS) $(BUILDDIR)/dirhtml 59 | @echo 60 | @echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml." 61 | 62 | singlehtml: 63 | $(SPHINXBUILD) -b singlehtml $(ALLSPHINXOPTS) $(BUILDDIR)/singlehtml 64 | @echo 65 | @echo "Build finished. The HTML page is in $(BUILDDIR)/singlehtml." 66 | 67 | pickle: 68 | $(SPHINXBUILD) -b pickle $(ALLSPHINXOPTS) $(BUILDDIR)/pickle 69 | @echo 70 | @echo "Build finished; now you can process the pickle files." 71 | 72 | json: 73 | $(SPHINXBUILD) -b json $(ALLSPHINXOPTS) $(BUILDDIR)/json 74 | @echo 75 | @echo "Build finished; now you can process the JSON files." 76 | 77 | htmlhelp: 78 | $(SPHINXBUILD) -b htmlhelp $(ALLSPHINXOPTS) $(BUILDDIR)/htmlhelp 79 | @echo 80 | @echo "Build finished; now you can run HTML Help Workshop with the" \ 81 | ".hhp project file in $(BUILDDIR)/htmlhelp." 82 | 83 | qthelp: 84 | $(SPHINXBUILD) -b qthelp $(ALLSPHINXOPTS) $(BUILDDIR)/qthelp 85 | @echo 86 | @echo "Build finished; now you can run "qcollectiongenerator" with the" \ 87 | ".qhcp project file in $(BUILDDIR)/qthelp, like this:" 88 | @echo "# qcollectiongenerator $(BUILDDIR)/qthelp/python-kdtree.qhcp" 89 | @echo "To view the help file:" 90 | @echo "# assistant -collectionFile $(BUILDDIR)/qthelp/python-kdtree.qhc" 91 | 92 | devhelp: 93 | $(SPHINXBUILD) -b devhelp $(ALLSPHINXOPTS) $(BUILDDIR)/devhelp 94 | @echo 95 | @echo "Build finished." 96 | @echo "To view the help file:" 97 | @echo "# mkdir -p $$HOME/.local/share/devhelp/python-kdtree" 98 | @echo "# ln -s $(BUILDDIR)/devhelp $$HOME/.local/share/devhelp/python-kdtree" 99 | @echo "# devhelp" 100 | 101 | epub: 102 | $(SPHINXBUILD) -b epub $(ALLSPHINXOPTS) $(BUILDDIR)/epub 103 | @echo 104 | @echo "Build finished. The epub file is in $(BUILDDIR)/epub." 105 | 106 | latex: 107 | $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex 108 | @echo 109 | @echo "Build finished; the LaTeX files are in $(BUILDDIR)/latex." 110 | @echo "Run \`make' in that directory to run these through (pdf)latex" \ 111 | "(use \`make latexpdf' here to do that automatically)." 112 | 113 | latexpdf: 114 | $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex 115 | @echo "Running LaTeX files through pdflatex..." 116 | $(MAKE) -C $(BUILDDIR)/latex all-pdf 117 | @echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex." 118 | 119 | latexpdfja: 120 | $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex 121 | @echo "Running LaTeX files through platex and dvipdfmx..." 122 | $(MAKE) -C $(BUILDDIR)/latex all-pdf-ja 123 | @echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex." 124 | 125 | text: 126 | $(SPHINXBUILD) -b text $(ALLSPHINXOPTS) $(BUILDDIR)/text 127 | @echo 128 | @echo "Build finished. The text files are in $(BUILDDIR)/text." 129 | 130 | man: 131 | $(SPHINXBUILD) -b man $(ALLSPHINXOPTS) $(BUILDDIR)/man 132 | @echo 133 | @echo "Build finished. The manual pages are in $(BUILDDIR)/man." 134 | 135 | texinfo: 136 | $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo 137 | @echo 138 | @echo "Build finished. The Texinfo files are in $(BUILDDIR)/texinfo." 139 | @echo "Run \`make' in that directory to run these through makeinfo" \ 140 | "(use \`make info' here to do that automatically)." 141 | 142 | info: 143 | $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo 144 | @echo "Running Texinfo files through makeinfo..." 145 | make -C $(BUILDDIR)/texinfo info 146 | @echo "makeinfo finished; the Info files are in $(BUILDDIR)/texinfo." 147 | 148 | gettext: 149 | $(SPHINXBUILD) -b gettext $(I18NSPHINXOPTS) $(BUILDDIR)/locale 150 | @echo 151 | @echo "Build finished. The message catalogs are in $(BUILDDIR)/locale." 152 | 153 | changes: 154 | $(SPHINXBUILD) -b changes $(ALLSPHINXOPTS) $(BUILDDIR)/changes 155 | @echo 156 | @echo "The overview file is in $(BUILDDIR)/changes." 157 | 158 | linkcheck: 159 | $(SPHINXBUILD) -b linkcheck $(ALLSPHINXOPTS) $(BUILDDIR)/linkcheck 160 | @echo 161 | @echo "Link check complete; look for any errors in the above output " \ 162 | "or in $(BUILDDIR)/linkcheck/output.txt." 163 | 164 | doctest: 165 | $(SPHINXBUILD) -b doctest $(ALLSPHINXOPTS) $(BUILDDIR)/doctest 166 | @echo "Testing of doctests in the sources finished, look at the " \ 167 | "results in $(BUILDDIR)/doctest/output.txt." 168 | 169 | xml: 170 | $(SPHINXBUILD) -b xml $(ALLSPHINXOPTS) $(BUILDDIR)/xml 171 | @echo 172 | @echo "Build finished. The XML files are in $(BUILDDIR)/xml." 173 | 174 | pseudoxml: 175 | $(SPHINXBUILD) -b pseudoxml $(ALLSPHINXOPTS) $(BUILDDIR)/pseudoxml 176 | @echo 177 | @echo "Build finished. The pseudo-XML files are in $(BUILDDIR)/pseudoxml." 178 | -------------------------------------------------------------------------------- /doc/conf.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # 3 | # python-kdtree documentation build configuration file, created by 4 | # sphinx-quickstart on Sat Mar 22 15:55:01 2014. 5 | # 6 | # This file is execfile()d with the current directory set to its 7 | # containing dir. 8 | # 9 | # Note that not all possible configuration values are present in this 10 | # autogenerated file. 11 | # 12 | # All configuration values have a default; values that are commented out 13 | # serve to show the default. 14 | 15 | import sys 16 | import os 17 | 18 | # If extensions (or modules to document with autodoc) are in another directory, 19 | # add these directories to sys.path here. If the directory is relative to the 20 | # documentation root, use os.path.abspath to make it absolute, like shown here. 21 | sys.path.insert(0, os.path.join(os.path.abspath('.'), '..')) 22 | 23 | import kdtree 24 | 25 | # -- General configuration ------------------------------------------------ 26 | 27 | # If your documentation needs a minimal Sphinx version, state it here. 28 | #needs_sphinx = '1.0' 29 | 30 | # Add any Sphinx extension module names here, as strings. They can be 31 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 32 | # ones. 33 | extensions = [] 34 | 35 | # Add any paths that contain templates here, relative to this directory. 36 | templates_path = ['_templates'] 37 | 38 | # The suffix of source filenames. 39 | source_suffix = '.rst' 40 | 41 | # The encoding of source files. 42 | #source_encoding = 'utf-8-sig' 43 | 44 | # The master toctree document. 45 | master_doc = 'index' 46 | 47 | # General information about the project. 48 | project = u'python-kdtree' 49 | copyright = kdtree.__author__ 50 | 51 | # The version info for the project you're documenting, acts as replacement for 52 | # |version| and |release|, also used in various other places throughout the 53 | # built documents. 54 | # 55 | # The short X.Y version. 56 | version = kdtree.__version__ 57 | # The full version, including alpha/beta/rc tags. 58 | release = kdtree.__version__ 59 | 60 | # The language for content autogenerated by Sphinx. Refer to documentation 61 | # for a list of supported languages. 62 | #language = None 63 | 64 | # There are two options for replacing |today|: either, you set today to some 65 | # non-false value, then it is used: 66 | #today = '' 67 | # Else, today_fmt is used as the format for a strftime call. 68 | #today_fmt = '%B %d, %Y' 69 | 70 | # List of patterns, relative to source directory, that match files and 71 | # directories to ignore when looking for source files. 72 | exclude_patterns = ['_build'] 73 | 74 | # The reST default role (used for this markup: `text`) to use for all 75 | # documents. 76 | #default_role = None 77 | 78 | # If true, '()' will be appended to :func: etc. cross-reference text. 79 | #add_function_parentheses = True 80 | 81 | # If true, the current module name will be prepended to all description 82 | # unit titles (such as .. function::). 83 | #add_module_names = True 84 | 85 | # If true, sectionauthor and moduleauthor directives will be shown in the 86 | # output. They are ignored by default. 87 | #show_authors = False 88 | 89 | # The name of the Pygments (syntax highlighting) style to use. 90 | pygments_style = 'sphinx' 91 | 92 | # A list of ignored prefixes for module index sorting. 93 | #modindex_common_prefix = [] 94 | 95 | # If true, keep warnings as "system message" paragraphs in the built documents. 96 | #keep_warnings = False 97 | 98 | 99 | # -- Options for HTML output ---------------------------------------------- 100 | 101 | # The theme to use for HTML and HTML Help pages. See the documentation for 102 | # a list of builtin themes. 103 | html_theme = 'default' 104 | 105 | # Theme options are theme-specific and customize the look and feel of a theme 106 | # further. For a list of options available for each theme, see the 107 | # documentation. 108 | #html_theme_options = {} 109 | 110 | # Add any paths that contain custom themes here, relative to this directory. 111 | #html_theme_path = [] 112 | 113 | # The name for this set of Sphinx documents. If None, it defaults to 114 | # " v documentation". 115 | #html_title = None 116 | 117 | # A shorter title for the navigation bar. Default is the same as html_title. 118 | #html_short_title = None 119 | 120 | # The name of an image file (relative to this directory) to place at the top 121 | # of the sidebar. 122 | #html_logo = None 123 | 124 | # The name of an image file (within the static path) to use as favicon of the 125 | # docs. This file should be a Windows icon file (.ico) being 16x16 or 32x32 126 | # pixels large. 127 | #html_favicon = None 128 | 129 | # Add any paths that contain custom static files (such as style sheets) here, 130 | # relative to this directory. They are copied after the builtin static files, 131 | # so a file named "default.css" will overwrite the builtin "default.css". 132 | html_static_path = ['_static'] 133 | 134 | # Add any extra paths that contain custom files (such as robots.txt or 135 | # .htaccess) here, relative to this directory. These files are copied 136 | # directly to the root of the documentation. 137 | #html_extra_path = [] 138 | 139 | # If not '', a 'Last updated on:' timestamp is inserted at every page bottom, 140 | # using the given strftime format. 141 | #html_last_updated_fmt = '%b %d, %Y' 142 | 143 | # If true, SmartyPants will be used to convert quotes and dashes to 144 | # typographically correct entities. 145 | #html_use_smartypants = True 146 | 147 | # Custom sidebar templates, maps document names to template names. 148 | #html_sidebars = {} 149 | 150 | # Additional templates that should be rendered to pages, maps page names to 151 | # template names. 152 | #html_additional_pages = {} 153 | 154 | # If false, no module index is generated. 155 | #html_domain_indices = True 156 | 157 | # If false, no index is generated. 158 | #html_use_index = True 159 | 160 | # If true, the index is split into individual pages for each letter. 161 | #html_split_index = False 162 | 163 | # If true, links to the reST sources are added to the pages. 164 | #html_show_sourcelink = True 165 | 166 | # If true, "Created using Sphinx" is shown in the HTML footer. Default is True. 167 | #html_show_sphinx = True 168 | 169 | # If true, "(C) Copyright ..." is shown in the HTML footer. Default is True. 170 | #html_show_copyright = True 171 | 172 | # If true, an OpenSearch description file will be output, and all pages will 173 | # contain a tag referring to it. The value of this option must be the 174 | # base URL from which the finished HTML is served. 175 | #html_use_opensearch = '' 176 | 177 | # This is the file name suffix for HTML files (e.g. ".xhtml"). 178 | #html_file_suffix = None 179 | 180 | # Output file base name for HTML help builder. 181 | htmlhelp_basename = 'python-kdtreedoc' 182 | 183 | 184 | # -- Options for LaTeX output --------------------------------------------- 185 | 186 | latex_elements = { 187 | # The paper size ('letterpaper' or 'a4paper'). 188 | #'papersize': 'letterpaper', 189 | 190 | # The font size ('10pt', '11pt' or '12pt'). 191 | #'pointsize': '10pt', 192 | 193 | # Additional stuff for the LaTeX preamble. 194 | #'preamble': '', 195 | } 196 | 197 | # Grouping the document tree into LaTeX files. List of tuples 198 | # (source start file, target name, title, 199 | # author, documentclass [howto, manual, or own class]). 200 | latex_documents = [ 201 | ('index', 'python-kdtree.tex', u'python-kdtree Documentation', 202 | u'Stefan Kögl', 'manual'), 203 | ] 204 | 205 | # The name of an image file (relative to this directory) to place at the top of 206 | # the title page. 207 | #latex_logo = None 208 | 209 | # For "manual" documents, if this is true, then toplevel headings are parts, 210 | # not chapters. 211 | #latex_use_parts = False 212 | 213 | # If true, show page references after internal links. 214 | #latex_show_pagerefs = False 215 | 216 | # If true, show URL addresses after external links. 217 | #latex_show_urls = False 218 | 219 | # Documents to append as an appendix to all manuals. 220 | #latex_appendices = [] 221 | 222 | # If false, no module index is generated. 223 | #latex_domain_indices = True 224 | 225 | 226 | # -- Options for manual page output --------------------------------------- 227 | 228 | # One entry per manual page. List of tuples 229 | # (source start file, name, description, authors, manual section). 230 | man_pages = [ 231 | ('index', 'python-kdtree', u'python-kdtree Documentation', 232 | [u'Stefan Kögl'], 1) 233 | ] 234 | 235 | # If true, show URL addresses after external links. 236 | #man_show_urls = False 237 | 238 | 239 | # -- Options for Texinfo output ------------------------------------------- 240 | 241 | # Grouping the document tree into Texinfo files. List of tuples 242 | # (source start file, target name, title, author, 243 | # dir menu entry, description, category) 244 | texinfo_documents = [ 245 | ('index', 'python-kdtree', u'python-kdtree Documentation', 246 | u'Stefan Kögl', 'python-kdtree', 'One line description of project.', 247 | 'Miscellaneous'), 248 | ] 249 | 250 | # Documents to append as an appendix to all manuals. 251 | #texinfo_appendices = [] 252 | 253 | # If false, no module index is generated. 254 | #texinfo_domain_indices = True 255 | 256 | # How to display URL addresses: 'footnote', 'no', or 'inline'. 257 | #texinfo_show_urls = 'footnote' 258 | 259 | # If true, do not generate a @detailmenu in the "Top" node's menu. 260 | #texinfo_no_detailmenu = False 261 | -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | from __future__ import absolute_import 4 | 5 | import sys 6 | import random 7 | import logging 8 | import unittest 9 | import doctest 10 | import collections 11 | from itertools import islice 12 | from bounded_priority_queue import BoundedPriorityQueue 13 | 14 | try: 15 | import coverage 16 | coverage.erase() 17 | coverage.start() 18 | 19 | except ImportError: 20 | coverage = None 21 | 22 | # import after starting coverage, to ensure that import-time code is covered 23 | import kdtree 24 | 25 | class RemoveTest(unittest.TestCase): 26 | 27 | 28 | def test_remove_duplicates(self): 29 | """ creates a tree with only duplicate points, and removes them all """ 30 | 31 | points = [(1,1)] * 100 32 | tree = kdtree.create(points) 33 | self.assertTrue(tree.is_valid()) 34 | 35 | random.shuffle(points) 36 | while points: 37 | point = points.pop(0) 38 | 39 | tree = tree.remove(point) 40 | 41 | # Check if the Tree is valid after the removal 42 | self.assertTrue(tree.is_valid()) 43 | 44 | # Check if the removal reduced the number of nodes by 1 (not more, not less) 45 | remaining_points = len(points) 46 | nodes_in_tree = len(list(tree.inorder())) 47 | self.assertEqual(nodes_in_tree, remaining_points) 48 | 49 | 50 | def test_remove(self, num=100): 51 | """ Tests random removal from a tree, multiple times """ 52 | 53 | for i in range(num): 54 | self.do_random_remove() 55 | 56 | 57 | def do_random_remove(self): 58 | """ Creates a random tree, removes all points in random order """ 59 | 60 | points = list(set(islice(random_points(), 0, 20))) 61 | tree = kdtree.create(points) 62 | self.assertTrue(tree.is_valid()) 63 | 64 | random.shuffle(points) 65 | while points: 66 | point = points.pop(0) 67 | 68 | tree = tree.remove(point) 69 | 70 | # Check if the Tree is valid after the removal 71 | self.assertTrue(tree.is_valid()) 72 | 73 | # Check if the point has actually been removed 74 | self.assertTrue(point not in [n.data for n in tree.inorder()]) 75 | 76 | # Check if the removal reduced the number of nodes by 1 (not more, not less) 77 | remaining_points = len(points) 78 | nodes_in_tree = len(list(tree.inorder())) 79 | self.assertEqual(nodes_in_tree, remaining_points) 80 | 81 | def test_remove_empty_tree(self): 82 | tree = kdtree.create(dimensions=2) 83 | tree.remove( (1, 2) ) 84 | self.assertFalse(bool(tree)) 85 | 86 | 87 | class AddTest(unittest.TestCase): 88 | 89 | def test_add(self, num=10): 90 | """ Tests random additions to a tree, multiple times """ 91 | 92 | for i in range(num): 93 | self.do_random_add() 94 | 95 | 96 | def do_random_add(self, num_points=100): 97 | 98 | points = list(set(islice(random_points(), 0, num_points))) 99 | tree = kdtree.create(dimensions=len(points[0])) 100 | for n, point in enumerate(points, 1): 101 | 102 | tree.add(point) 103 | 104 | self.assertTrue(tree.is_valid()) 105 | 106 | self.assertTrue(point in [node.data for node in tree.inorder()]) 107 | 108 | nodes_in_tree = len(list(tree.inorder())) 109 | self.assertEqual(nodes_in_tree, n) 110 | 111 | 112 | class InvalidTreeTests(unittest.TestCase): 113 | 114 | 115 | def test_invalid_child(self): 116 | """ Children on wrong subtree invalidate Tree """ 117 | child = kdtree.KDNode( (3, 2) ) 118 | child.axis = 2 119 | tree = kdtree.create([(2, 3)]) 120 | tree.left=child 121 | self.assertFalse(tree.is_valid()) 122 | 123 | tree = kdtree.create([(4, 1)]) 124 | tree.right=child 125 | self.assertFalse(tree.is_valid()) 126 | 127 | 128 | def test_different_dimensions(self): 129 | """ Can't create Tree for Points of different dimensions """ 130 | points = [ (1, 2), (2, 3, 4) ] 131 | self.assertRaises(ValueError, kdtree.create, points) 132 | 133 | 134 | class TreeTraversals(unittest.TestCase): 135 | 136 | def test_same_length(self): 137 | tree = random_tree() 138 | 139 | inorder_len = len(list(tree.inorder())) 140 | preorder_len = len(list(tree.preorder())) 141 | postorder_len = len(list(tree.postorder())) 142 | 143 | self.assertEqual(inorder_len, preorder_len) 144 | self.assertEqual(preorder_len, postorder_len) 145 | 146 | 147 | 148 | class BalanceTests(unittest.TestCase): 149 | 150 | 151 | def test_rebalance(self): 152 | 153 | tree = random_tree(1) 154 | while tree.is_balanced: 155 | tree.add(random_point()) 156 | 157 | tree = tree.rebalance() 158 | self.assertTrue(tree.is_balanced) 159 | 160 | 161 | 162 | class NearestNeighbor(unittest.TestCase): 163 | 164 | def test_search_knn(self): 165 | points = [(50, 20), (51, 19), (1, 80)] 166 | tree = kdtree.create(points) 167 | point = (48, 18) 168 | 169 | all_dist = [] 170 | for p in tree.inorder(): 171 | dist = p.dist(point) 172 | all_dist.append([p, dist]) 173 | 174 | all_dist = sorted(all_dist, key = lambda n:n[1]) 175 | 176 | result = tree.search_knn(point, 1) 177 | self.assertEqual(result[0][1], all_dist[0][1]) 178 | 179 | result = tree.search_knn(point, 2) 180 | self.assertEqual(result[0][1], all_dist[0][1]) 181 | self.assertEqual(result[1][1], all_dist[1][1]) 182 | 183 | result = tree.search_knn(point, 3) 184 | self.assertEqual(result[0][1], all_dist[0][1]) 185 | self.assertEqual(result[1][1], all_dist[1][1]) 186 | self.assertEqual(result[2][1], all_dist[2][1]) 187 | 188 | def test_search_nn(self, nodes=100): 189 | points = list(islice(random_points(), 0, nodes)) 190 | tree = kdtree.create(points) 191 | point = random_point() 192 | 193 | nn, dist = tree.search_nn(point) 194 | best, best_dist = self.find_best(tree, point) 195 | self.assertEqual(best_dist, dist, msg=', '.join(repr(p) for p in points) + ' / ' + repr(point)) 196 | 197 | 198 | def test_search_nn2(self): 199 | points = [(1,2,3),(5,1,2),(9,3,4),(3,9,1),(4,8,3),(9,1,1),(5,0,0), 200 | (1,1,1),(7,2,2),(5,9,1),(1,1,9),(9,8,7),(2,3,4),(4,5,4.01)] 201 | tree = kdtree.create(points) 202 | point = (2,5,6) 203 | 204 | nn, dist = tree.search_nn(point) 205 | best, best_dist = self.find_best(tree, point) 206 | self.assertEqual(best_dist, dist) 207 | 208 | 209 | def test_search_nn3(self): 210 | points = [(0, 25, 73), (1, 91, 85), (1, 47, 12), (2, 90, 20), 211 | (2, 66, 79), (2, 46, 27), (4, 48, 99), (5, 73, 64), (7, 42, 70), 212 | (7, 34, 60), (8, 86, 80), (10, 27, 14), (15, 64, 39), (17, 74, 24), 213 | (18, 58, 12), (18, 58, 5), (19, 14, 2), (20, 88, 11), (20, 28, 58), 214 | (20, 79, 48), (21, 32, 8), (21, 46, 41), (22, 6, 4), (22, 42, 68), 215 | (22, 62, 42), (24, 70, 96), (27, 77, 57), (27, 47, 39), (28, 61, 19), 216 | (30, 28, 22), (34, 13, 85), (34, 39, 96), (34, 90, 32), (39, 7, 45), 217 | (40, 61, 53), (40, 69, 50), (41, 45, 16), (41, 15, 44), (42, 40, 19), 218 | (45, 6, 68), (46, 79, 91), (47, 91, 86), (47, 50, 24), (48, 57, 64), 219 | (49, 21, 72), (49, 87, 21), (49, 41, 62), (54, 94, 32), (56, 14, 54), 220 | (56, 93, 2), (58, 34, 44), (58, 27, 42), (59, 62, 80), (60, 69, 69), 221 | (61, 67, 35), (62, 31, 50), (63, 9, 93), (63, 46, 95), (64, 31, 2), 222 | (64, 2, 36), (65, 23, 96), (66, 94, 69), (67, 98, 10), (67, 40, 88), 223 | (68, 4, 15), (68, 1, 6), (68, 88, 72), (70, 24, 53), (70, 31, 87), 224 | (71, 95, 26), (74, 80, 34), (75, 59, 99), (75, 15, 25), (76, 90, 99), 225 | (77, 75, 19), (77, 68, 26), (80, 19, 98), (82, 90, 50), (82, 87, 37), 226 | (84, 88, 59), (85, 76, 61), (85, 89, 20), (85, 64, 64), (86, 55, 92), 227 | (86, 15, 69), (87, 48, 46), (87, 67, 47), (89, 81, 65), (89, 87, 39), 228 | (89, 87, 3), (91, 65, 87), (94, 37, 74), (94, 20, 92), (95, 95, 49), 229 | (96, 15, 80), (96, 27, 39), (97, 87, 32), (97, 43, 7), (98, 78, 10), 230 | (99, 64, 55)] 231 | 232 | tree = kdtree.create(points) 233 | point = (66, 54, 29) 234 | 235 | nn, dist = tree.search_nn(point) 236 | best, best_dist = self.find_best(tree, point) 237 | self.assertEqual(best_dist, dist) 238 | 239 | 240 | 241 | def find_best(self, tree, point): 242 | best = None 243 | best_dist = None 244 | for p in tree.inorder(): 245 | dist = p.dist(point) 246 | if best is None or dist < best_dist: 247 | best = p 248 | best_dist = dist 249 | return best, best_dist 250 | 251 | def test_search_nn_dist(self): 252 | """ tests search_nn_dist() according to bug #8 """ 253 | 254 | points = [(x,y) for x in range(10) for y in range(10)] 255 | tree = kdtree.create(points) 256 | nn = tree.search_nn_dist((5,5), 2.5) 257 | 258 | self.assertEquals(len(nn), 4) 259 | self.assertTrue( (6,6) in nn) 260 | self.assertTrue( (5,5) in nn) 261 | self.assertTrue( (5,6) in nn) 262 | self.assertTrue( (6,5) in nn) 263 | 264 | 265 | def test_search_nn_dist_random(self): 266 | 267 | for n in range(50): 268 | tree = random_tree() 269 | point = random_point() 270 | points = tree.inorder() 271 | 272 | points = sorted(points, key=lambda p: p.dist(point)) 273 | 274 | for p in points: 275 | dist = p.dist(point) 276 | nn = tree.search_nn_dist(point, dist) 277 | 278 | for pn in points: 279 | if pn in nn: 280 | self.assertTrue(pn.dist(point) < dist, '%s in %s but %s < %s' % (pn, nn, pn.dist(point), dist)) 281 | else: 282 | self.assertTrue(pn.dist(point) >= dist, '%s not in %s but %s >= %s' % (pn, nn, pn.dist(point), dist)) 283 | 284 | 285 | class PointTypeTests(unittest.TestCase): 286 | """ test using different types as points """ 287 | 288 | def test_point_types(self): 289 | emptyTree = kdtree.create(dimensions=3) 290 | point1 = (2, 3, 4) 291 | point2 = [4, 5, 6] 292 | Point = collections.namedtuple('Point', 'x y z') 293 | point3 = Point(5, 3, 2) 294 | tree = kdtree.create([point1, point2, point3]) 295 | res, dist = tree.search_nn( (1, 2, 3) ) 296 | 297 | self.assertEqual(res, kdtree.KDNode( (2, 3, 4) )) 298 | 299 | 300 | class PayloadTests(unittest.TestCase): 301 | """ test tree.add() with payload """ 302 | 303 | def test_payload(self, nodes=100, dimensions=3): 304 | points = list(islice(random_points(dimensions=dimensions), 0, nodes)) 305 | tree = kdtree.create(dimensions=dimensions) 306 | 307 | for i, p in enumerate(points): 308 | tree.add(p).payload = i 309 | 310 | for i, p in enumerate(points): 311 | self.assertEqual(i, tree.search_nn(p)[0].payload) 312 | 313 | 314 | class BoundedPriorityQueueTests(unittest.TestCase): 315 | """ Test bounded priority queue """ 316 | 317 | def test_size(self): 318 | bpq = self.get_test_bpq() 319 | self.assertEqual(bpq.size(), 5) 320 | 321 | def test_max(self): 322 | bpq = self.get_test_bpq() 323 | # note: not 4.6! The bpq keeps the smallest k values (5 in this test). 324 | self.assertEqual(bpq.max(), 3.2) 325 | 326 | def test_extract(self): 327 | bpq = self.get_test_bpq() 328 | nodes = self.get_test_nodes() 329 | nodes = sorted(nodes, key=lambda n: n[1], reverse=True) 330 | # get rid of the 2 largest nodes, since the bpq only keeps the 331 | # 5 smallest of the 7 nodes 332 | nodes = nodes[2:] 333 | for i in range(5): 334 | node, dist = bpq.extract_max() 335 | exp_node, exp_dist = nodes.pop(0) 336 | self.assertEqual(dist, exp_dist) 337 | self.assertEqual(bpq.size(), 0) 338 | 339 | def get_test_bpq(self): 340 | bound = 5 341 | bpq = BoundedPriorityQueue(bound) 342 | for n in self.get_test_nodes(): 343 | bpq.add(n) 344 | return bpq 345 | 346 | def get_test_nodes(self): 347 | return [ 348 | (None, 0.1), 349 | (None, 0.25), 350 | (None, 1.33), 351 | (None, 3.2), 352 | (None, 4.6), 353 | (None, 0.4), 354 | (None, 4.0) 355 | ] 356 | 357 | 358 | def random_tree(nodes=20, dimensions=3, minval=0, maxval=100): 359 | points = list(islice(random_points(), 0, nodes)) 360 | tree = kdtree.create(points) 361 | return tree 362 | 363 | def random_point(dimensions=3, minval=0, maxval=100): 364 | return tuple(random.randint(minval, maxval) for _ in range(dimensions)) 365 | 366 | def random_points(dimensions=3, minval=0, maxval=100): 367 | while True: 368 | yield random_point(dimensions, minval, maxval) 369 | 370 | 371 | if __name__ == '__main__': 372 | 373 | suite = unittest.TestLoader().loadTestsFromModule(sys.modules[__name__]) 374 | suite.addTest(doctest.DocTestSuite(kdtree)) 375 | runner = unittest.TextTestRunner(verbosity=2) 376 | result = runner.run(suite) 377 | 378 | if not result.wasSuccessful(): 379 | sys.exit(1) 380 | 381 | if coverage is not None: 382 | coverage.stop() 383 | coverage.report([kdtree]) 384 | coverage.erase() 385 | -------------------------------------------------------------------------------- /kdtree.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | 4 | """A Python implemntation of a kd-tree 5 | 6 | This package provides a simple implementation of a kd-tree in Python. 7 | https://en.wikipedia.org/wiki/K-d_tree 8 | """ 9 | 10 | from __future__ import print_function 11 | 12 | import operator 13 | import math 14 | from collections import deque 15 | from functools import wraps 16 | from bounded_priority_queue import BoundedPriorityQueue 17 | 18 | __author__ = u'Stefan Kögl ' 19 | __version__ = '0.15' 20 | __website__ = 'https://github.com/stefankoegl/kdtree' 21 | __license__ = 'ISC license' 22 | 23 | 24 | class Node(object): 25 | """ A Node in a kd-tree 26 | 27 | A tree is represented by its root node, and every node represents 28 | its subtree""" 29 | 30 | def __init__(self, data=None, left=None, right=None): 31 | self.data = data 32 | self.left = left 33 | self.right = right 34 | 35 | 36 | @property 37 | def is_leaf(self): 38 | """ Returns True if a Node has no subnodes 39 | 40 | >>> Node().is_leaf 41 | True 42 | 43 | >>> Node( 1, left=Node(2) ).is_leaf 44 | False 45 | """ 46 | return (not self.data) or \ 47 | (all(not bool(c) for c, p in self.children)) 48 | 49 | 50 | def preorder(self): 51 | """ iterator for nodes: root, left, right """ 52 | 53 | if not self: 54 | return 55 | 56 | yield self 57 | 58 | if self.left: 59 | for x in self.left.preorder(): 60 | yield x 61 | 62 | if self.right: 63 | for x in self.right.preorder(): 64 | yield x 65 | 66 | 67 | def inorder(self): 68 | """ iterator for nodes: left, root, right """ 69 | 70 | if not self: 71 | return 72 | 73 | if self.left: 74 | for x in self.left.inorder(): 75 | yield x 76 | 77 | yield self 78 | 79 | if self.right: 80 | for x in self.right.inorder(): 81 | yield x 82 | 83 | 84 | def postorder(self): 85 | """ iterator for nodes: left, right, root """ 86 | 87 | if not self: 88 | return 89 | 90 | if self.left: 91 | for x in self.left.postorder(): 92 | yield x 93 | 94 | if self.right: 95 | for x in self.right.postorder(): 96 | yield x 97 | 98 | yield self 99 | 100 | 101 | @property 102 | def children(self): 103 | """ 104 | Returns an iterator for the non-empty children of the Node 105 | 106 | The children are returned as (Node, pos) tuples where pos is 0 for the 107 | left subnode and 1 for the right. 108 | 109 | >>> len(list(create(dimensions=2).children)) 110 | 0 111 | 112 | >>> len(list(create([ (1, 2) ]).children)) 113 | 0 114 | 115 | >>> len(list(create([ (2, 2), (2, 1), (2, 3) ]).children)) 116 | 2 117 | """ 118 | 119 | if self.left and self.left.data is not None: 120 | yield self.left, 0 121 | if self.right and self.right.data is not None: 122 | yield self.right, 1 123 | 124 | 125 | def set_child(self, index, child): 126 | """ Sets one of the node's children 127 | 128 | index 0 refers to the left, 1 to the right child """ 129 | 130 | if index == 0: 131 | self.left = child 132 | else: 133 | self.right = child 134 | 135 | 136 | def height(self): 137 | """ 138 | Returns height of the (sub)tree, without considering 139 | empty leaf-nodes 140 | 141 | >>> create(dimensions=2).height() 142 | 0 143 | 144 | >>> create([ (1, 2) ]).height() 145 | 1 146 | 147 | >>> create([ (1, 2), (2, 3) ]).height() 148 | 2 149 | """ 150 | 151 | min_height = int(bool(self)) 152 | return max([min_height] + [c.height()+1 for c, p in self.children]) 153 | 154 | 155 | def get_child_pos(self, child): 156 | """ Returns the position if the given child 157 | 158 | If the given node is the left child, 0 is returned. If its the right 159 | child, 1 is returned. Otherwise None """ 160 | 161 | for c, pos in self.children: 162 | if child == c: 163 | return pos 164 | 165 | 166 | def __repr__(self): 167 | return '<%(cls)s - %(data)s>' % \ 168 | dict(cls=self.__class__.__name__, data=repr(self.data)) 169 | 170 | 171 | def __nonzero__(self): 172 | return self.data is not None 173 | 174 | __bool__ = __nonzero__ 175 | 176 | def __eq__(self, other): 177 | if isinstance(other, tuple): 178 | return self.data == other 179 | else: 180 | return self.data == other.data 181 | 182 | def __hash__(self): 183 | return id(self) 184 | 185 | 186 | def require_axis(f): 187 | """ Check if the object of the function has axis and sel_axis members """ 188 | 189 | @wraps(f) 190 | def _wrapper(self, *args, **kwargs): 191 | if None in (self.axis, self.sel_axis): 192 | raise ValueError('%(func_name) requires the node %(node)s ' 193 | 'to have an axis and a sel_axis function' % 194 | dict(func_name=f.__name__, node=repr(self))) 195 | 196 | return f(self, *args, **kwargs) 197 | 198 | return _wrapper 199 | 200 | 201 | 202 | class KDNode(Node): 203 | """ A Node that contains kd-tree specific data and methods """ 204 | 205 | 206 | def __init__(self, data=None, left=None, right=None, axis=None, 207 | sel_axis=None, dimensions=None): 208 | """ Creates a new node for a kd-tree 209 | 210 | If the node will be used within a tree, the axis and the sel_axis 211 | function should be supplied. 212 | 213 | sel_axis(axis) is used when creating subnodes of the current node. It 214 | receives the axis of the parent node and returns the axis of the child 215 | node. """ 216 | super(KDNode, self).__init__(data, left, right) 217 | self.axis = axis 218 | self.sel_axis = sel_axis 219 | self.dimensions = dimensions 220 | 221 | 222 | @require_axis 223 | def add(self, point): 224 | """ 225 | Adds a point to the current node or iteratively 226 | descends to one of its children. 227 | 228 | Users should call add() only to the topmost tree. 229 | """ 230 | 231 | current = self 232 | while True: 233 | check_dimensionality([point], dimensions=current.dimensions) 234 | 235 | # Adding has hit an empty leaf-node, add here 236 | if current.data is None: 237 | current.data = point 238 | return current 239 | 240 | # split on self.axis, recurse either left or right 241 | if point[current.axis] < current.data[current.axis]: 242 | if current.left is None: 243 | current.left = current.create_subnode(point) 244 | return current.left 245 | else: 246 | current = current.left 247 | else: 248 | if current.right is None: 249 | current.right = current.create_subnode(point) 250 | return current.right 251 | else: 252 | current = current.right 253 | 254 | 255 | @require_axis 256 | def create_subnode(self, data): 257 | """ Creates a subnode for the current node """ 258 | 259 | return self.__class__(data, 260 | axis=self.sel_axis(self.axis), 261 | sel_axis=self.sel_axis, 262 | dimensions=self.dimensions) 263 | 264 | 265 | @require_axis 266 | def find_replacement(self): 267 | """ Finds a replacement for the current node 268 | 269 | The replacement is returned as a 270 | (replacement-node, replacements-parent-node) tuple """ 271 | 272 | if self.right: 273 | child, parent = self.right.extreme_child(min, self.axis) 274 | else: 275 | child, parent = self.left.extreme_child(max, self.axis) 276 | 277 | return (child, parent if parent is not None else self) 278 | 279 | 280 | def should_remove(self, point, node): 281 | """ checks if self's point (and maybe identity) matches """ 282 | if not self.data == point: 283 | return False 284 | 285 | return (node is None) or (node is self) 286 | 287 | 288 | @require_axis 289 | def remove(self, point, node=None): 290 | """ Removes the node with the given point from the tree 291 | 292 | Returns the new root node of the (sub)tree. 293 | 294 | If there are multiple points matching "point", only one is removed. The 295 | optional "node" parameter is used for checking the identity, once the 296 | removeal candidate is decided.""" 297 | 298 | # Recursion has reached an empty leaf node, nothing here to delete 299 | if not self: 300 | return 301 | 302 | # Recursion has reached the node to be deleted 303 | if self.should_remove(point, node): 304 | return self._remove(point) 305 | 306 | # Remove direct subnode 307 | if self.left and self.left.should_remove(point, node): 308 | self.left = self.left._remove(point) 309 | 310 | elif self.right and self.right.should_remove(point, node): 311 | self.right = self.right._remove(point) 312 | 313 | # Recurse to subtrees 314 | if point[self.axis] <= self.data[self.axis]: 315 | if self.left: 316 | self.left = self.left.remove(point, node) 317 | 318 | if point[self.axis] >= self.data[self.axis]: 319 | if self.right: 320 | self.right = self.right.remove(point, node) 321 | 322 | return self 323 | 324 | 325 | @require_axis 326 | def _remove(self, point): 327 | # we have reached the node to be deleted here 328 | 329 | # deleting a leaf node is trivial 330 | if self.is_leaf: 331 | self.data = None 332 | return self 333 | 334 | # we have to delete a non-leaf node here 335 | 336 | # find a replacement for the node (will be the new subtree-root) 337 | root, max_p = self.find_replacement() 338 | 339 | # self and root swap positions 340 | tmp_l, tmp_r = self.left, self.right 341 | self.left, self.right = root.left, root.right 342 | root.left, root.right = tmp_l if tmp_l is not root else self, tmp_r if tmp_r is not root else self 343 | self.axis, root.axis = root.axis, self.axis 344 | 345 | # Special-case if we have not chosen a direct child as the replacement 346 | if max_p is not self: 347 | pos = max_p.get_child_pos(root) 348 | max_p.set_child(pos, self) 349 | max_p.remove(point, self) 350 | 351 | else: 352 | root.remove(point, self) 353 | 354 | return root 355 | 356 | 357 | @property 358 | def is_balanced(self): 359 | """ Returns True if the (sub)tree is balanced 360 | 361 | The tree is balanced if the heights of both subtrees differ at most by 362 | 1 """ 363 | 364 | left_height = self.left.height() if self.left else 0 365 | right_height = self.right.height() if self.right else 0 366 | 367 | if abs(left_height - right_height) > 1: 368 | return False 369 | 370 | return all(c.is_balanced for c, _ in self.children) 371 | 372 | 373 | def rebalance(self): 374 | """ 375 | Returns the (possibly new) root of the rebalanced tree 376 | """ 377 | 378 | return create([x.data for x in self.inorder()]) 379 | 380 | 381 | def axis_dist(self, point, axis): 382 | """ 383 | Squared distance at the given axis between 384 | the current Node and the given point 385 | """ 386 | return math.pow(self.data[axis] - point[axis], 2) 387 | 388 | 389 | def dist(self, point): 390 | """ 391 | Squared distance between the current Node 392 | and the given point 393 | """ 394 | r = range(self.dimensions) 395 | return sum([self.axis_dist(point, i) for i in r]) 396 | 397 | 398 | def search_knn(self, point, k, dist=None): 399 | """ Return the k nearest neighbors of point and their distances 400 | 401 | point must be an actual point, not a node. 402 | 403 | k is the number of results to return. The actual results can be less 404 | (if there aren't more nodes to return) or more in case of equal 405 | distances. 406 | 407 | dist is a distance function, expecting two points and returning a 408 | distance value. Distance values can be any compareable type. 409 | 410 | The result is an ordered list of (node, distance) tuples. 411 | """ 412 | 413 | if dist is None: 414 | get_dist = lambda n: n.dist(point) 415 | else: 416 | get_dist = lambda n: dist(n.data, point) 417 | 418 | results = BoundedPriorityQueue(k) 419 | 420 | self._search_node(point, k, results, get_dist) 421 | 422 | # We sort the final result by the distance in the tuple 423 | # (, distance) 424 | BY_VALUE = lambda kv: kv[1] 425 | return sorted(results.items(), key=BY_VALUE) 426 | 427 | 428 | def _search_node(self, point, k, results, get_dist): 429 | if not self: 430 | return 431 | 432 | nodeDist = get_dist(self) 433 | 434 | # Add current node to the priority queue if it closer than 435 | # at least one point in the queue. This functionality is 436 | # taken care of by BoundedPriorityQueue. 437 | results.add((self, nodeDist)) 438 | 439 | # get the splitting plane 440 | split_plane = self.data[self.axis] 441 | # get the squared distance between the point and the splitting plane 442 | # (squared since all distances are squared). 443 | plane_dist = point[self.axis] - split_plane 444 | plane_dist2 = plane_dist * plane_dist 445 | 446 | # Search the side of the splitting plane that the point is in 447 | if point[self.axis] < split_plane: 448 | if self.left is not None: 449 | self.left._search_node(point, k, results, get_dist) 450 | else: 451 | if self.right is not None: 452 | self.right._search_node(point, k, results, get_dist) 453 | 454 | # Search the other side of the splitting plane if it may contain 455 | # points closer than the farthest point in the current results. 456 | if plane_dist2 < results.max() or results.size() < k: 457 | if point[self.axis] < self.data[self.axis]: 458 | if self.right is not None: 459 | self.right._search_node(point, k, results, get_dist) 460 | else: 461 | if self.left is not None: 462 | self.left._search_node(point, k, results, get_dist) 463 | 464 | 465 | @require_axis 466 | def search_nn(self, point, dist=None): 467 | """ 468 | Search the nearest node of the given point 469 | 470 | point must be an actual point, not a node. The nearest node to the 471 | point is returned. If a location of an actual node is used, the Node 472 | with this location will be returned (not its neighbor). 473 | 474 | dist is a distance function, expecting two points and returning a 475 | distance value. Distance values can be any compareable type. 476 | 477 | The result is a (node, distance) tuple. 478 | """ 479 | 480 | return next(iter(self.search_knn(point, 1, dist)), None) 481 | 482 | 483 | @require_axis 484 | def search_nn_dist(self, point, distance, best=None): 485 | """ 486 | Search the n nearest nodes of the given point which are within given 487 | distance 488 | 489 | point must be a location, not a node. A list containing the n nearest 490 | nodes to the point within the distance will be returned. 491 | """ 492 | 493 | if best is None: 494 | best = [] 495 | 496 | # consider the current node 497 | if self.dist(point) < distance: 498 | best.append(self) 499 | 500 | # sort the children, nearer one first (is this really necessairy?) 501 | children = sorted(self.children, key=lambda c_p1: c_p1[0].dist(point)) 502 | 503 | for child, p in children: 504 | # check if child node needs to be recursed 505 | if self.axis_dist(point, self.axis) < math.pow(distance, 2): 506 | child.search_nn_dist(point, distance, best) 507 | 508 | return best 509 | 510 | 511 | @require_axis 512 | def is_valid(self): 513 | """ Checks recursively if the tree is valid 514 | 515 | It is valid if each node splits correctly """ 516 | 517 | if not self: 518 | return True 519 | 520 | if self.left and self.data[self.axis] < self.left.data[self.axis]: 521 | return False 522 | 523 | if self.right and self.data[self.axis] > self.right.data[self.axis]: 524 | return False 525 | 526 | return all(c.is_valid() for c, _ in self.children) or self.is_leaf 527 | 528 | 529 | def extreme_child(self, sel_func, axis): 530 | """ Returns a child of the subtree and its parent 531 | 532 | The child is selected by sel_func which is either min or max 533 | (or a different function with similar semantics). """ 534 | 535 | max_key = lambda child_parent: child_parent[0].data[axis] 536 | 537 | 538 | # we don't know our parent, so we include None 539 | me = [(self, None)] if self else [] 540 | 541 | child_max = [c.extreme_child(sel_func, axis) for c, _ in self.children] 542 | # insert self for unknown parents 543 | child_max = [(c, p if p is not None else self) for c, p in child_max] 544 | 545 | candidates = me + child_max 546 | 547 | if not candidates: 548 | return None, None 549 | 550 | return sel_func(candidates, key=max_key) 551 | 552 | 553 | 554 | def create(point_list=None, dimensions=None, axis=0, sel_axis=None): 555 | """ Creates a kd-tree from a list of points 556 | 557 | All points in the list must be of the same dimensionality. 558 | 559 | If no point_list is given, an empty tree is created. The number of 560 | dimensions has to be given instead. 561 | 562 | If both a point_list and dimensions are given, the numbers must agree. 563 | 564 | Axis is the axis on which the root-node should split. 565 | 566 | sel_axis(axis) is used when creating subnodes of a node. It receives the 567 | axis of the parent node and returns the axis of the child node. """ 568 | 569 | if not point_list and not dimensions: 570 | raise ValueError('either point_list or dimensions must be provided') 571 | 572 | elif point_list: 573 | dimensions = check_dimensionality(point_list, dimensions) 574 | 575 | # by default cycle through the axis 576 | sel_axis = sel_axis or (lambda prev_axis: (prev_axis+1) % dimensions) 577 | 578 | if not point_list: 579 | return KDNode(sel_axis=sel_axis, axis=axis, dimensions=dimensions) 580 | 581 | # Sort point list and choose median as pivot element 582 | point_list = list(point_list) 583 | point_list.sort(key=lambda point: point[axis]) 584 | median = len(point_list) // 2 585 | 586 | loc = point_list[median] 587 | left = create(point_list[:median], dimensions, sel_axis(axis)) 588 | right = create(point_list[median + 1:], dimensions, sel_axis(axis)) 589 | return KDNode(loc, left, right, axis=axis, sel_axis=sel_axis, dimensions=dimensions) 590 | 591 | 592 | def check_dimensionality(point_list, dimensions=None): 593 | dimensions = dimensions or len(point_list[0]) 594 | for p in point_list: 595 | if len(p) != dimensions: 596 | raise ValueError('All Points in the point_list must have the same dimensionality') 597 | 598 | return dimensions 599 | 600 | 601 | 602 | def level_order(tree, include_all=False): 603 | """ Returns an iterator over the tree in level-order 604 | 605 | If include_all is set to True, empty parts of the tree are filled 606 | with dummy entries and the iterator becomes infinite. """ 607 | 608 | q = deque() 609 | q.append(tree) 610 | while q: 611 | node = q.popleft() 612 | yield node 613 | 614 | if include_all or node.left: 615 | q.append(node.left or node.__class__()) 616 | 617 | if include_all or node.right: 618 | q.append(node.right or node.__class__()) 619 | 620 | 621 | 622 | def visualize(tree, max_level=100, node_width=10, left_padding=5): 623 | """ Prints the tree to stdout """ 624 | 625 | height = min(max_level, tree.height()-1) 626 | max_width = pow(2, height) 627 | 628 | per_level = 1 629 | in_level = 0 630 | level = 0 631 | 632 | for node in level_order(tree, include_all=True): 633 | 634 | if in_level == 0: 635 | print() 636 | print() 637 | print(' '*left_padding, end=' ') 638 | 639 | width = int(max_width*node_width/per_level) 640 | 641 | node_str = (str(node.data) if node else '').center(width) 642 | print(node_str, end=' ') 643 | 644 | in_level += 1 645 | 646 | if in_level == per_level: 647 | in_level = 0 648 | per_level *= 2 649 | level += 1 650 | 651 | if level > height: 652 | break 653 | 654 | print() 655 | print() 656 | --------------------------------------------------------------------------------