【python】通行网格地图四叉树化

受到Leecode 427题的启发，427. 建立四叉树
想将由0和1组成的网格地图绘制为四叉树地图，0表示可通行网格，1表示不可通行网格。

import matplotlib.pyplot as plt  
import matplotlib.patches as patches  
import numpy as np  
from matplotlib import colors  
  
  
class Node:  
    def __init__(self, x0, y0, w, h, data):  
        self.x0 = x0  
        self.y0 = y0  
        self.width = w  
        self.height = h  
        self.data = data  
        self.children = []  
        self.passable = np.sum(data == 1) / (w * h) <= 0.9  # 障碍物比例<= 0.2 可通行  
  
  
class QuadTree:  
    def __init__(self, x, y, w, h, data):  
        self.root = Node(x, y, w, h, data)  
        self.root.passable = True  
        if self.root.passable:  
            self.subdivide(self.root)  
  
    def subdivide(self, node):  
		x, y, w, h = node.x0, node.y0, node.width // 2, node.height // 2
		# Divide the node's data into quadrants
        data_quadrants = [  
            node.data[0:h, 0:w],  
            node.data[h: node.height, 0:w],  
            node.data[0:h, w:node.width],  
            node.data[h:node.height, w:node.width]  
        ]  
  
        for i, data in enumerate(data_quadrants):  
            width = w  
            height = h  
            if (i // 2):  
                width = node.width - w  
            if (i % 2):  
                height = node.height - h  
            child = Node(x + (i // 2) * w, y + (i % 2) * h, width, height, data)  
            node.children.append(child)  
            if len(np.unique(data)) > 1 and w > 1 and h > 1:  
                self.subdivide(child)  
  
    def draw(self, ax):  
        def _draw(node):  
            color = 'black' if not node.passable else 'white'  
            # ax.add_patch(patches.Rectangle((node.x0, node.y0), node.width, node.height, facecolor=color, edgecolor='black'))  
            ax.add_patch(  
                patches.Rectangle((node.x0, node.y0), node.width, node.height, facecolor=color, edgecolor='black'))  
            for child in node.children:  
                if child is not None:  
                    _draw(child)  
  
        _draw(self.root)  
  
  
    def print_leaves(self, node=None):  
        if node is None:  
            node = self.root  
        if all(child is None for child in node.children):  
            print(  
                f"Leaf Node: x={node.x0}, y={node.y0}, width={node.width}, height={node.height}, passable={node.passable}")  
        else:  
            for child in node.children:  
                if child is not None:  
                    self.print_leaves(child)  
  
  
# Create a 32x32 grid map with random obstacles  
map_data = np.loadtxt("./map_txt/Random_32x32_.2.txt")  
  
width, height = map_data.shape  
  
# Create a quadtree  
qt = QuadTree(0, 0, width, height, map_data)  
  
cmap = colors.ListedColormap(['white', 'black', ])  
  
# Draw the original map  
fig, axs = plt.subplots(1, 2, figsize=(10, 5))  
axs[0].imshow(map_data, cmap=cmap, origin='lower')  
axs[0].set_title('Original Map')  
  
# Draw the quadtree  
axs[1].set_xlim(0, width)  
axs[1].set_ylim(0, height)  
qt.draw(axs[1])  
axs[1].set_title('Quadtree Map')  
  
plt.show()

运行结果