当前位置：首页 > news >正文

四叉树实现四边形网格

news 2025/6/7 1:12:41

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
# 四叉树节点
class QuadNode:def __init__(self, x, y, width, height, depth):self.x = xself.y = yself.width = widthself.height = heightself.depth = depthself.children = []self.id = None  # 单元编号def is_leaf(self):return len(self.children) == 0# 四叉树
class Quadtree:def __init__(self, x, y, width, height, max_depth):self.root = QuadNode(x, y, width, height, 0)self.max_depth = max_depthdef subdivide(self, node):if node.depth >= self.max_depth:returnw, h = node.width / 2, node.height / 2d = node.depth + 1node.children = [QuadNode(node.x,       node.y,       w, h, d),QuadNode(node.x + w,   node.y,       w, h, d),QuadNode(node.x,       node.y + h,   w, h, d),QuadNode(node.x + w,   node.y + h,   w, h, d)]for child in node.children:self.subdivide(child)def get_leaf_nodes(self, node=None):if node is None:node = self.rootif node.is_leaf():return [node]else:leaves = []for child in node.children:leaves.extend(self.get_leaf_nodes(child))return leaves# 编号 & 绘图函数
def draw_quadtree_with_numbering(quadtree):fig, ax = plt.subplots(figsize=(8, 8))leaves = quadtree.get_leaf_nodes()# 单元编号排序：先按y排，再按x排（自底向上、左到右）leaves.sort(key=lambda n: (n.y, n.x))# 给单元编号for i, node in enumerate(leaves, 1):node.id = i# 统计所有节点坐标，去重node_coords = set()for node in leaves:corners = [(node.x, node.y),(node.x + node.width, node.y),(node.x, node.y + node.height),(node.x + node.width, node.y + node.height)]node_coords.update(corners)# 节点编号排序：按 y, x 排序（自底向上、左到右）sorted_nodes = sorted(list(node_coords), key=lambda p: (p[1], p[0]))node_id_map = {coord: idx+1 for idx, coord in enumerate(sorted_nodes)}# 画单元格和编号for node in leaves:rect = patches.Rectangle((node.x, node.y), node.width, node.height,linewidth=1, edgecolor='black', facecolor='none')ax.add_patch(rect)# 单元编号（圆圈 + 数字）center_x = node.x + node.width / 2center_y = node.y + node.height / 2circle = patches.Circle((center_x, center_y), radius=node.width * 0.08,edgecolor='blue', facecolor='lightyellow', linewidth=1)ax.add_patch(circle)ax.text(center_x, center_y, str(node.id), fontsize=8, color='red', ha='center', va='center')# 画节点及编号for coord in sorted_nodes:ax.plot(coord[0], coord[1], 'ko', markersize=3)ax.text(coord[0], coord[1], str(node_id_map[coord]), fontsize=7, color='green',ha='center', va='center', bbox=dict(boxstyle="circle,pad=0.15", fc="white", ec="green", lw=0.8))ax.set_xticks([])ax.set_yticks([])ax.set_xticklabels([])ax.set_yticklabels([])ax.axis('off')plt.show()# 生成单元编号，顺序为左下、左上、右上、右下（逆时针）# 生成单元顶点编号数组，顺序为逆时针elements = []for node in leaves:ll = (node.x, node.y)  # 左下ul = (node.x, node.y + node.height)  # 左上ur = (node.x + node.width, node.y + node.height)  # 右上lr = (node.x + node.width, node.y)  # 右下cell_ids = [node_id_map[ll], node_id_map[lr], node_id_map[ur], node_id_map[ul]]elements.append(cell_ids)elements = np.array(elements)  # NC x 4nodes = np.array(sorted_nodes)  # NN x 2return elements, nodesif __name__ == "__main__":tree = Quadtree(0, 0, 1, 1, max_depth=1)tree.subdivide(tree.root)elements, nodes = draw_quadtree_with_numbering(tree)print(f"单元数量: {len(elements)}")print(f"节点数量: {len(nodes)}")# 打印前5个单元和节点print("前5个单元:")for e in elements[:5]:print(e)print("前5个节点:")for n in nodes[:5]:print(n)