quad_tree.go

package main

import (
	"fmt"
	"math"
)

type PetrolPump struct {
	location	Location
	name string
}

type Location struct {
	x int
	y int
}

type Region struct {
	topLeftlocation Location
	width int
	height int

	pumps []PetrolPump

	northWestRegion *Region
	northEastRegion *Region
	southWestRegion *Region
	southEastRegion *Region
}

type QuadTree struct {
	rootRegion *Region
	bucketSize int // max number of pumps allowed in one region
}

func NewQuadTree(x, y, width, height, bucketSize int) *QuadTree {
	rootRegion := &Region{
		topLeftlocation: Location{x, y},
		width: width,
		height: height,
	}
	return &QuadTree{rootRegion, bucketSize}
}

func (q *QuadTree) locationWithinRegion(location Location, region *Region) bool {
	if (location.x >= region.topLeftlocation.x &&
		location.x < region.topLeftlocation.x + region.width &&
		location.y <= region.topLeftlocation.y &&
		location.y > region.topLeftlocation.y - region.height) {
		return true
	}
	return false
}

func (q *QuadTree) isLeafRegion(region *Region) bool {
	if (region.northWestRegion == nil && region.northEastRegion == nil &&
		region.southEastRegion == nil && region.southWestRegion == nil) {
		return true
	}
	return false
}

func (q *QuadTree) divideRegion(region *Region) {
	x := region.topLeftlocation.x
	y := region.topLeftlocation.y
	halfWidth := (region.width / 2)
	halfHeight := (region.height / 2) 
	midX := x + halfWidth
	midY := y - halfHeight

	region.northWestRegion = &Region{
		topLeftlocation: Location{x, y},
		width: halfWidth,
		height: halfHeight,
	}
	region.northEastRegion = &Region{
		topLeftlocation: Location{midX, y},
		width: region.width - halfWidth,
		height: halfHeight,
	}
	region.southWestRegion = &Region{
		topLeftlocation: Location{x, midY},
		width: halfWidth,
		height: region.height - halfHeight,
	}
	region.southEastRegion = &Region{
		topLeftlocation: Location{midX, midY},
		width: region.width - halfWidth,
		height: region.height - halfHeight,
	}
}

func (q *QuadTree) insertNewRegion(region *Region, pump PetrolPump) {
	if !q.locationWithinRegion(pump.location, region) {
		return
	}

	if !q.isLeafRegion(region) {
		if region.northWestRegion != nil && q.locationWithinRegion(pump.location, region.northWestRegion) {
			q.insertNewRegion(region.northWestRegion, pump)
		}
		if region.northEastRegion != nil && q.locationWithinRegion(pump.location, region.northEastRegion) {
			q.insertNewRegion(region.northEastRegion, pump)
		}
		if region.southWestRegion != nil && q.locationWithinRegion(pump.location, region.southWestRegion) {
			q.insertNewRegion(region.southWestRegion, pump)
		}
		if region.southEastRegion != nil && q.locationWithinRegion(pump.location, region.southEastRegion) {
			q.insertNewRegion(region.southEastRegion, pump)
		}
		return
	}

	if len(region.pumps) < q.bucketSize {
		region.pumps = append(region.pumps, pump)
		return
	}

	q.divideRegion(region)
	region.pumps = append(region.pumps, pump)
	for _, pump := range region.pumps {
		q.insertNewRegion(region.northWestRegion, pump)
		q.insertNewRegion(region.northEastRegion, pump)
		q.insertNewRegion(region.southWestRegion, pump)
		q.insertNewRegion(region.southEastRegion, pump)
	}
	region.pumps = []PetrolPump{}
}

func (q *QuadTree) Insert(pump PetrolPump) {
	q.insertNewRegion(q.rootRegion, pump)
}

func (q *QuadTree) getPumps(region *Region, location Location, pumps *[]PetrolPump) {
	if !q.locationWithinRegion(location, region) {
		return
	}

	if q.isLeafRegion(region) {
		for _, pump := range region.pumps {
			*pumps = append(*pumps, pump)
		}
		return
	}

	if region.northWestRegion != nil && q.locationWithinRegion(location, region.northWestRegion) {
		q.getPumps(region.northWestRegion, location, pumps)
		return
	}
	if region.northEastRegion != nil && q.locationWithinRegion(location, region.northEastRegion) {
		q.getPumps(region.northEastRegion, location, pumps)
		return
	}
	if region.southEastRegion != nil && q.locationWithinRegion(location, region.southEastRegion) {
		q.getPumps(region.southEastRegion, location, pumps)
		return
	}
	if region.southWestRegion != nil && q.locationWithinRegion(location, region.southWestRegion) {
		q.getPumps(region.southWestRegion, location, pumps)
		return
	}
}

func (q *QuadTree) GetPumpsInMyRegion(location Location) []PetrolPump {
	pumps := []PetrolPump{}
	q.getPumps(q.rootRegion, location, &pumps)
	return pumps
}

func (q *QuadTree) locationWithinRadius(location, center Location, radius int) bool {
	dist := math.Sqrt(math.Pow(float64(location.x - center.x), 2) + math.Pow(float64(location.y - center.y), 2))
	return dist <= float64(radius)
}

func (q *QuadTree) regionWithinRadius(region *Region, center Location, radius int) bool {
	locations := []Location{
		region.topLeftlocation,
		Location{region.topLeftlocation.x + region.width - 1, region.topLeftlocation.y},
		Location{region.topLeftlocation.x, region.topLeftlocation.y - region.height + 1},
		Location{region.topLeftlocation.x + region.width - 1, region.topLeftlocation.y - region.height + 1},
	}
	for _, location := range locations {
		if q.locationWithinRadius(location, center, radius) {
			return true
		}
	}
	return false
}

func (q *QuadTree) shouldProbeRegion(region *Region, center Location, radius int) bool {
	return q.regionWithinRadius(region, center, radius) || q.locationWithinRegion(center, region)
}

func (q *QuadTree) getPumpsWithinRadius(region *Region, center Location, radius int, pumps *[]PetrolPump) {
	if q.shouldProbeRegion(region, center, radius) && q.isLeafRegion(region) {
		for _, pump := range region.pumps {
			if q.locationWithinRadius(pump.location, center, radius) {
				*pumps = append(*pumps, pump)
			}	
		}
	}

	if region.northWestRegion != nil && q.shouldProbeRegion(region.northWestRegion, center, radius) {
		q.getPumpsWithinRadius(region.northWestRegion, center, radius, pumps)
	}
	if region.northEastRegion != nil && q.shouldProbeRegion(region.northEastRegion, center, radius) {
		q.getPumpsWithinRadius(region.northEastRegion, center, radius, pumps)
	}
	if region.southWestRegion != nil && q.shouldProbeRegion(region.southWestRegion, center, radius) {
		q.getPumpsWithinRadius(region.southWestRegion, center, radius, pumps)
	}
	if region.southEastRegion != nil && q.shouldProbeRegion(region.southEastRegion, center, radius) {
		q.getPumpsWithinRadius(region.southEastRegion, center, radius, pumps)
	}
}

func (q *QuadTree) getAllPumpsInRadius(center Location, radius int) []PetrolPump {
	var pumps []PetrolPump
	q.getPumpsWithinRadius(q.rootRegion, center, radius, &pumps)
	return pumps
}

func main() {
	qt := NewQuadTree(2, 9, 5, 6, 5)
	pumps := []PetrolPump{
		PetrolPump{Location{3, 8}, "a"},
		PetrolPump{Location{5, 7}, "b"},
		PetrolPump{Location{6, 9}, "c"},
		PetrolPump{Location{6, 7}, "d"},
		PetrolPump{Location{5, 5}, "e"},
		PetrolPump{Location{2, 5}, "f"},
	}
	for _, pump := range pumps {
		qt.Insert(pump)
	}

	fmt.Println(qt.GetPumpsInMyRegion(Location{2, 4}))
	fmt.Println(qt.GetPumpsInMyRegion(Location{5, 8}))
	fmt.Println(qt.GetPumpsInMyRegion(Location{8, 8}))
	fmt.Println(qt.GetPumpsInMyRegion(Location{2, 6}))
	fmt.Println(qt.getAllPumpsInRadius(Location{4, 7}, 0))
	fmt.Println(qt.getAllPumpsInRadius(Location{4, 7}, 1))
	fmt.Println(qt.getAllPumpsInRadius(Location{4, 7}, 2))
	fmt.Println(qt.getAllPumpsInRadius(Location{4, 7}, 3))
	
	fmt.Println()
	qt = NewQuadTree(0, 7, 8, 8, 1)
	pumps = []PetrolPump{
		PetrolPump{Location{1, 6}, "a"},
		PetrolPump{Location{4, 5}, "b"},
		PetrolPump{Location{2, 4}, "c"},
		PetrolPump{Location{6, 3}, "d"},
		PetrolPump{Location{5, 2}, "e"},
		PetrolPump{Location{3, 2}, "f"},
		PetrolPump{Location{0, 2}, "h"},
		PetrolPump{Location{4, 0}, "i"},
	}
	for _, pump := range pumps {
		qt.Insert(pump)
	}

	fmt.Println(qt.GetPumpsInMyRegion(Location{2, 4}))
	fmt.Println(qt.GetPumpsInMyRegion(Location{5, 7}))
	fmt.Println(qt.GetPumpsInMyRegion(Location{7, 7}))
	fmt.Println(qt.GetPumpsInMyRegion(Location{2, 6}))
	fmt.Println(qt.GetPumpsInMyRegion(Location{0, 2}))
	fmt.Println(qt.getAllPumpsInRadius(Location{6, 2}, 0))
	fmt.Println(qt.getAllPumpsInRadius(Location{6, 2}, 1))
	fmt.Println(qt.getAllPumpsInRadius(Location{6, 2}, 2))
	fmt.Println(qt.getAllPumpsInRadius(Location{6, 2}, 3))
	fmt.Println(qt.getAllPumpsInRadius(Location{6, 2}, 4))
	fmt.Println(qt.getAllPumpsInRadius(Location{6, 2}, 5))
	fmt.Println(qt.getAllPumpsInRadius(Location{6, 2}, 6))
	fmt.Println(qt.getAllPumpsInRadius(Location{6, 2}, 7))
}