一份测试kdtree的python代码

最新推荐文章于 2024-01-27 10:38:47 发布

原创最新推荐文章于 2024-01-27 10:38:47 发布 · 355 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python #kdtree

Python 专栏收录该内容

0 篇文章

订阅专栏

本文提供了一份使用Python实现K-D树的代码示例。通过测试代码，作者展示了如何构造K-D树并使用它来处理多维空间数据。该实现基于wiki百科上的K-D树定义，适合用于项目中涉及多维数据搜索和组织的场景。

一份测试kdtree的python代码：

忙完再补~~

自己的项目里要用到kdtree，先测试了一下子。直接贴代码，更多信息详见对应的wiki百科(https://en.wikipedia.org/wiki/K-d_tree)：

# ////Example implementation////
# srcUrl: https://en.wikipedia.org/wiki/K-d_tree

from collections import namedtuple
from operator import itemgetter
from pprint import pformat


class Node(namedtuple('Node', 'location left_child right_child')):
    def __repr__(self):
        return pformat(tuple(self))


def kdtree(point_list, depth: int = 0):
    if not point_list:
        return None

    k = len(point_list[0])  # assumes all points have the same dimension
    # Select axis based on depth so that axis cycles through all valid values
    axis = depth % k

    # Sort point list by axis and choose median as pivot element
    point_list.sort(key=itemgetter(axis))
    median = len(point_list) // 2

    # Create node and construct subtrees
    return Node(
        location=point_list[median],
        left_child=kdtree(point_list[:median], depth + 1),
        right_child=kdtree(point_list[median + 1:], depth + 1)
    )


def main():
    """Example usage"""
    point_list = [(7, 2), (5, 4), (9, 6), (4, 7), (8, 1), (2, 3)]
    tree = kdtree(point_list)
    print(tree)


if __name__ == '__main__':
    main()