쿼드트리

	import random as random
	import matplotlib.pyplot as plt
	import matplotlib.patches as patches
	import time
	class point:
	def __init__(self, x, y):
	self.x = x
	self.y = y
	class node:
	def __init__(self, bodies, width, height, posX, posY):
	self.hasleaf = False
	self.width = width
	self.height = height
	self.bodies = bodies
	self.posX = posX
	self.posY = posY
	self.leaf = []
	def generate_leaf(self, depth):
	q1 = []
	q2 = []
	q3 = []
	q4 = []
	for point in self.bodies:
	if point.x < (self.posX + (self.width / 2.0)):
	if point.y < (self.posY + (self.height / 2.0)):
	q1.append(point)
	else:
	q3.append(point)
	else:
	if point.y < (self.posY + (self.height / 2.0)):
	q2.append(point)
	else:
	q4.append(point)
	if len(q1) >= 1 and depth >= 1:
	self.hasleaf = True
	k1 = node(q1, self.width / 2.0, self.height / 2.0, self.posX, self.posY)
	k1.generate_leaf(depth - 1)
	self.leaf.append(k1)
	if len(q2) >= 1 and depth >= 1:
	self.hasleaf = True
	k2 = node(q2, self.width / 2.0, self.height / 2.0, self.posX + (self.width / 2.0), self.posY)
	k2.generate_leaf(depth - 1)
	self.leaf.append(k2)
	if len(q3) >= 1 and depth >= 1:
	self.hasleaf = True
	k3 = node(q3, self.width / 2.0, self.height / 2.0, self.posX, self.posY + (self.height / 2.0))
	k3.generate_leaf(depth - 1)
	self.leaf.append(k3)
	if len(q4) >= 1 and depth >= 1:
	self.hasleaf = True
	k4 = node(q4, self.width / 2.0, self.height / 2.0, self.posX + (self.width / 2.0), self.posY + (self.height / 2.0))
	k4.generate_leaf(depth - 1)
	self.leaf.append(k4)
	def postorder(node):
	if not node.leaf:
	ret = []
	ret.append(node)
	return ret
	else:
	leaves = []
	for leaf in node.leaf:
	leaves += postorder(leaf)
	leaves.append(node)
	return leaves
	if __name__ == "__main__":
	p = []
	for x in range(200):
	p.append(point(random.uniform(0.0, 10.0), random.uniform(0.0, 10.0)))
	tree = node(p, 10.0, 10.0, 0.0 ,0.0)
	tree.generate_leaf(10)
	print(len(tree.leaf[0].bodies))
	print(len(tree.leaf[1].bodies))
	print(len(tree.leaf[2].bodies))
	print(len(tree.leaf[3].bodies))
	q = postorder(tree)
	for x in q:
	print(x.width)
	fig = plt.figure(figsize = (12,8))
	ax = fig.add_subplot(111)
	for n in q:
	ax.add_patch(patches.Rectangle((n.posX, n.posY), n.width, n.height, fill=False))
	x = [point.x for point in p]
	y = [point.y for point in p]
	plt.plot(x, y, 'ro')
	plt.show()

view raw quadtree.py hosted with ❤ by GitHub

Barnes-Hut 천체 시뮬레이션 만들고 있는데 필요해서 만들어봄.
이진 트리하고 비슷하게 만들면 되는데 이진 트리보다 구분해야 하는 차원의 개수가 하나 늘어서 만들 때 꼬일 수 있음.