git.fiddlerwoaroof.com
Raw Blame History
from __future__ import print_function
import collections

#
#  (0,0)  (1,0) | (2,0) (3,0)
#  (0,1)  (1,1) | (2,1) (3,1)
#  -------------+------------
#  (0,2)  (1,2) | (2,2) (3,2)
#  (0,3)  (1,3) | (2,3) (3,3)
#

class Rect(object):
	@property
	def width(self): return self.size[0]
	@property
	def height(self): return self.size[1]

	@property
	def x(self): return self.pos[0]
	@property
	def y(self): return self.pos[1]

	@classmethod
	def random_rect(cls, xrange, yrange, dimrange):
		x = random.randrange(*xrange)
		y = random.randrange(*yrange)
		size = map(lambda x: random.randrange(*x), dimrange)
		return cls(x,y, *size)

	def __repr__(self):
		return 'Rect(%d, %d, %d, %d)' % (self.pos + self.size)

	def fill(self, array, value):
		x,y = self.pos
		w,h = self.size
		array[x:x+w,y:y+h] = value

	def __init__(self, x,y, w,h):
		self.pos = (x,y)
		self.size = (w,h)

class QuadTree(object):
	MAX_OBJECTS = 10
	MAX_LEVELS = 8

	def __init__(self, pLevel, pBounds):
		self.level = pLevel
		self.objects = []
		self.bounds = pBounds
		self.nodes = [None, None, None, None]

	def clear(self):
		del self.objects[:]
		for idx,node in enumerate(self.nodes):
			if node is not None:
				node.clear()
				self.nodes[idx] = None

	def split(self):
		subWidth = self.bounds.size[0] / 2
		subHeight = self.bounds.size[1] / 2
		x, y = self.bounds.pos

		self.nodes[0] = QuadTree(self.level+1, Rect(x+subWidth, y, subWidth, subHeight))
		self.nodes[1] = QuadTree(self.level+1, Rect(x, y, subWidth, subHeight))
		self.nodes[2] = QuadTree(self.level+1, Rect(x, y + subHeight, subWidth, subHeight))
		self.nodes[3] = QuadTree(self.level+1, Rect(x + subWidth, y + subHeight, subWidth, subHeight))

	def getIndex(self, rect):
		index = -1
		verticalMidpoint = self.bounds.x + self.bounds.width/2
		horizontalMidpoint = self.bounds.y + self.bounds.width/2
		topquad = rect.y < horizontalMidpoint and rect.y + rect.height < horizontalMidpoint
		bottomquad = rect.y > horizontalMidpoint

		if rect.x < verticalMidpoint and rect.x + rect.width < verticalMidpoint:
			if topquad:
				index = 1
			elif bottomquad:
				index = 2
		elif rect.x > verticalMidpoint:
			if topquad:
				index = 0
			elif bottomquad:
				index = 3

		return index

	def insert(self, rect):
		if self.nodes[0] is not None:
			index = self.getIndex(rect)
			if index != -1:
				self.nodes[index].insert(rect)
				return

		self.objects.append(rect)

		if len(self.objects) > self.MAX_OBJECTS and self.level < self.MAX_LEVELS:
			if self.nodes[0] is None:
				self.split()

			i = 0
			while i < len(self.objects):
				index = self.getIndex(self.objects[i])
				if index != -1:
					self.nodes[index].insert(self.objects.pop(i))
				else:
					i += 1

	def retrieve(self, rect):
		result = []
		index = self.getIndex(rect)
		if index != -1 and self.nodes[0] is not None:
			result.extend(self.nodes[index].retrieve(rect))

		result.extend(self.objects)
		return result

	def locate_rect(self, rect):
		cur = self
		index = cur.getIndex(rect)

		while index != -1:
			yield index
			cur = self.nodes[index]




import random

objects = {Rect.random_rect( (1,37), (1,200), ( (2,10), (2,10) ) ) for ___ in range(50)}
quad = QuadTree(0, Rect(0,0, 48,211))
quad.clear()
for obj in objects:
	quad.insert(obj)

collided = set()
oot = set()

for obj in objects:
	out = quad.retrieve(obj)

	if obj in collided: continue
	oot.add(obj)

	for col in out:
		collided.add(col)
		print(obj, 'collides with', col)
	else:
		print('---')


import numpy
field = numpy.zeros( (48,211) )
for obj in objects:
	obj.fill(field, 2)
for obj in oot:
	obj.fill(field, 1)

for row in field:
	for cell in row:
		print('%d' % cell, end='')
	print()