commit:97b0a22977a4a5d66ca4695345bfa126909b5029
author:Chip Black
committer:Chip Black
date:Tue Jul 15 20:30:26 2008 -0500
parents:80a3486b45abc4d8c81da29580c21d402f152f6f
Added QuadTree implementation for storing rects
diff --git a/Richter/util.py b/Richter/util.py
line changes: +62/-0
index d3cc058..52ae449
--- a/Richter/util.py
+++ b/Richter/util.py
@@ -1,3 +1,6 @@
+import pygame
+from OpenGL.GL import *
+
 def n2ceil(n):
 	"Find the next power of two greater than or equal to n"
 	assert n > 0
@@ -8,3 +11,62 @@ def n2ceil(n):
 		x <<= 1
 	# FAIL
 	return 1024
+
+
+def splitrect(rect, point):
+	"Split a rect into four rects with point in the center"
+	assert rect.collidepoint(point)
+	return [
+		pygame.Rect(rect.topleft, (point[0] - rect.left, point[1] - rect.top)),
+		pygame.Rect(rect.midtop, (rect.right - point[0], point[1] - rect.top)),
+		pygame.Rect(rect.midleft, (point[0] - rect.left, rect.bottom - point[1])),
+		pygame.Rect(rect.center, (rect.right - point[0], rect.bottom - point[1]))
+	]
+
+
+class QuadTreeR:
+	"""Implements a QuadTree containing arbitrary rects
+	We number the quadrants like so:
+	    ^
+	  3 | 4
+	x --+-->
+	  1 | 2
+	    y
+	"""
+	def __init__(self, bbox):
+		self.bbox = pygame.Rect(bbox)
+		self.rects = []
+		self.subregions = None
+
+
+	def add(self, rect, obj):
+		if not self.subregions:
+			self.rects.append((pygame.Rect(rect), obj))
+			self.subregions = [QuadTreeR(r) for r in splitrect(self.bbox, self.bbox.center)]
+		else:
+			for s in self.subregions:
+				if s.bbox.contains(rect):
+					s.add(rect, obj)
+					break
+			else:
+				self.rects.append((pygame.Rect(rect), obj))
+	
+
+	def collideRectWalk(self, rect):
+		"""generator that walks through the QuadTree, returning rects
+		that collide with the given rect"""
+		for (r, o) in self.rects:
+			if r.colliderect(rect):
+				yield o
+		if self.subregions:
+			for s in self.subregions:
+				if s.bbox.contains(rect):
+					for o in s.collideRectWalk(rect):
+						yield o
+	
+
+	def collideRect(self, rect):
+		for o in self.collisionWalk(rect):
+			return o
+		else:
+			return None