可视化四叉树

四叉树(Q-Tree)是一种树形数据结构。四叉树的定义是:它的每个节点下至多可以有四个子节点,通常把一部分二维空间细分为四个象限或区域并把该区域里的相关信息存入到四叉树节点中。

更多原理请参考大佬写的原理讲解以及伪码: https://www.cnblogs.com/wellp/p/8536745.html

这里我们用opencv来绘制出四叉树对应的矩形区域,通过graphviz来可视化树的结构。鼠标左键双击可以添加数据点,右键双击则删除一个最近的数据点。
效果图:
在这里插入图片描述
树结构可视化:
说明:可视化这里因为为了省事用的现成的库,没有仔细研究,节点上标的是区域的左上角坐标,边上标的值意义不大。。
在这里插入图片描述
我们来看看代码实现:

import numpy as np
import cv2
from graphviz import Digraph

class QuadNode():
    def __init__(self, depth, x, y, width, height):
        self.rect = (x,y,width,height)
        self.data = []
        self.depth = depth
        self.sub = []

color = {0:(255,0,0), 1:(0,255,0), 2:(0,0,255), 3:(255,255,0), 4:(255,0,255), 5:(0,255,255), 6:(100,0,0)}
line_size = {0:6, 1:4, 2:3, 3:2, 4:1, 5:1, 6:1}

def traverse(root, img):
    tree_dic = {}
    if root:
        print(root.depth)
        if len(root.sub) > 0 or len(root.data) > 0:
            tree_dic[str((root.rect[0], root.rect[1]))] = {}
            for node in root.sub:
                t_dic = traverse(node, img)
                if len(t_dic) > 0:
                    tree_dic[str((root.rect[0],root.rect[1]))][str((node.rect[0], node.rect[1]))] = t_dic
            cv2.rectangle(img, (root.rect[0], root.rect[1]),
                      (root.rect[0] + root.rect[2], root.rect[1] + root.rect[3]),
                      color[root.depth],
                      line_size[root.depth])
    return tree_dic

def is_in_rect(rect, point):
    # 判断点point是否在矩形区域rect内
    if point[0] >= rect[0] and point[0] <= rect[0] + rect[2]\
            and point[1] >= rect[1] and point[1] <= rect[1] + rect[3]:
        return True
    else:
        return False

def create_child(root):
    x, y, width, height = root.rect
    c_depth = root.depth + 1
    # 左上的孩子
    tl = QuadNode(c_depth, x, y, width//2, height//2)
    # 右上的孩子
    tr = QuadNode(c_depth, x+width//2, y, width//2, height//2)
    # 左下的孩子
    dl = QuadNode(c_depth, x, y + width // 2, width // 2, height // 2)
    # 右下的孩子
    dr = QuadNode(c_depth, x + width // 2, y + height // 2, width // 2, height // 2)
    root.sub.append(tr)
    root.sub.append(tl)
    root.sub.append(dl)
    root.sub.append(dr)

def quad_insert(root, point):
    if len(root.sub) > 0:
        # 如果root有孩子,判断point应该放置于哪一个孩子节点
        for node in root.sub:
            if is_in_rect(node.rect, point):
                print('insert ',point, ' into', node.rect)
                quad_insert(node, point)
    elif len(root.data) > 0:
        # 如果root存储了一个对象,为root创建4个孩子
        create_child(root)
        # 将root中的对象移到对应的孩子节点
        data = root.data[0]
        root.data = []
        for node in root.sub:
            if is_in_rect(node.rect, data):
                quad_insert(root, data)
        # 划分区域判断point应该放置于哪一个孩子节点
        for node in root.sub:
            if is_in_rect(node.rect, point):
                quad_insert(node, point)
    elif len(root.data) == 0:
        root.data.append(point)

def add(event, x, y, flags, param):
    global tree_dic,img, points, root
    if event==cv2.EVENT_LBUTTONDBLCLK:
        point = (x,y)
        points.append(point)
        print(point)
        quad_insert(root, point)
        # 同时绘制出数据点
        cv2.circle(img, point, 1, (0, 0, 0), -1, lineType=cv2.LINE_AA)
        cv2.putText(img, str(point), point, cv2.FONT_ITALIC, 0.3, (0, 0, 0), 1)
        tree_dic = traverse(root, img)
    elif event==cv2.EVENT_RBUTTONDBLCLK:
        point = (x, y)
        if len(points) > 0:
            min_d = (points[0][0]-x)**2 + (points[0][1] - y)**2
            min = points[0]
            for p in points:
                d = (p[0]-x)**2 + (p[1] - y)**2
                if d < min_d:
                    min_d = d
                    min = p
            points.remove(min)
            print('r',point)
            img = np.zeros((400, 400, 3), np.uint8)
            img += 255
            root = QuadNode(0, 0, 0, 400, 400)
            for point in points:
                quad_insert(root, point)
                # 同时绘制出数据点
                cv2.circle(img, point, 1, (0,0,0), -1, lineType=cv2.LINE_AA)
                cv2.putText(img, str(point), point, cv2.FONT_ITALIC, 0.3, (0,0,0), 1)
            tree_dic = traverse(root, img)

def plot_model(tree, name):
    g = Digraph("G", filename=name, format='png', strict=False)
    first_label = list(tree.keys())[0]
    g.node("0", first_label)
    _sub_plot(g, tree, "0")
    g.view()


vroot = "0"


def _sub_plot(g, tree, inc):
    global vroot

    first_label = list(tree.keys())[0]
    ts = tree[first_label]
    for i in ts.keys():
        if isinstance(tree[first_label][i], dict):
            vroot = str(int(vroot) + 1)
            g.node(vroot, list(tree[first_label][i].keys())[0])
            g.edge(inc, vroot, str(i))
            _sub_plot(g, tree[first_label][i], vroot)
        else:
            root = str(int(vroot) + 1)
            g.node(vroot, tree[first_label][i])
            g.edge(inc, vroot, str(i))



points = []
img = np.zeros((400,400,3), np.uint8)
img += 255
# 创建空的四叉树
root = QuadNode(0,0,0,400,400)
tree_dic = {}

if __name__ == '__main__':

    cv2.namedWindow('quadtree')
    cv2.setMouseCallback('quadtree', add)
    while True:
        cv2.imshow('quadtree', img)
        key = cv2.waitKey(20)
        if key == ord('q'):
            # 可视化树结构
            print(tree_dic)
            plot_model(tree_dic, "guadtree.gv")
            break

    cv2.destroyAllWindows()
原文链接:加载失败,请重新获取