Quadtree Query - Python Pygame

import pygame
from pygame import gfxdraw
import random

# quadTree constants
CAPACITY = 4

# This breaks down the display area into populated rectangles. The population is based on points
class quadTree:
    def __init__(s, rect): # boundry is a rectangle
        s.rect = rect
        s.points = []
        s.divided = False

    def insert(s, point):
        global CAPACITY # maximum capacity on each quadrant

        # if the point inserted doesn't fit into this quadrant then don't add it to
        # the points list. Exit insert()!
        if not s.rect.collidepoint(int(point.x), int(point.y)): return False

        # check if the quadrant has room for more points
        if len(s.points) < CAPACITY:
            s.points.append(point) # if there's room add it to the points list
            # return True to inform parent quadtree the point has been added
            return True
        else:
            # if the quadrant is full then divide the quadrant into
            # four parts north west, north east, south west, south east...
            if not s.divided:
                s.subDivide()

            # ...and insert the point into one of the four new quadrants
            if s.nw.insert(point):
                return True
            elif s.ne.insert(point):
                return True
            elif s.sw.insert(point):
                return True
            elif s.se.insert(point):
                return True

    def subDivide(s):
        # create four new quadtrees which belong to the parent quadtree
        size = (s.rect.width / 2, s.rect.height / 2)
        s.nw = quadTree(pygame.Rect(s.rect.topleft, size))
        s.ne = quadTree(pygame.Rect(s.rect.midtop, size))
        s.sw = quadTree(pygame.Rect(s.rect.midleft, size))
        s.se = quadTree(pygame.Rect(s.rect.center, size))
        s.divided = True # the parent quadtree can no longer be divided

    def query(s, rect, pointsArray = []):
        # if the queried rect doesn't intersect with the quadrant
        # exit query
        if not s.rect.colliderect(rect): return

        # if it does then append points to the pointsArray
        for p in s.points:
            if rect.collidepoint(p): pointsArray.append(p)

        # because the quadtree is recursive we need to check its children
        # to see if they're within the query parameters (rect)
        if s.divided:
            s.nw.query(rect, pointsArray)
            s.ne.query(rect, pointsArray)
            s.sw.query(rect, pointsArray)
            s.se.query(rect, pointsArray)

    def show(s):
        global PD # Primary Display

        if not s.points: return

        # draw all the points in this quadrant to the Primary Display
        for p in s.points:
            pygame.draw.circle(PD, (255, 0, 0), (int(p.x), int(p.y)), 3)

        # if the quadrant is divided then:
        # 1. draw the division lines and then
        # 2. get its children to draw their points and division lines
        if s.divided:
            # draw the boundry to help visualise what's going on
            pygame.gfxdraw.hline(PD, int(s.rect.x), int(s.rect.right), int(s.rect.centery), (255, 255, 255))
            pygame.gfxdraw.vline(PD, int(s.rect.centerx), int(s.rect.y), int(s.rect.bottom), (255, 255, 255))

            s.nw.show()
            s.ne.show()
            s.sw.show()
            s.se.show()

# pygame display settings
DR = pygame.Rect((0, 0, 800, 600)) # Display Rectangle
HDW, HDH = DR.center # H = half

# set up pygame
pygame.init()
PD = pygame.display.set_mode(DR.size) # primary display based of the size of Display Rect (800, 600)

# create a quadtree instance and initialise it with the display rectange (size of the display)
QT = quadTree(DR)

# previous position of the mouse pointer
px, py = -1, -1

# press ESC to exit
while True:
    # process events
    e = pygame.event.get()
    if pygame.key.get_pressed()[pygame.K_ESCAPE]: break

    # get position of the mouse
    mx, my = pygame.mouse.get_pos()


    # get the status of the mouse buttons
    mb = pygame.mouse.get_pressed()

    # if left mouse button has been pressed then add randomised points
    if mb[0]:
        # if current mouse position is different to the previous mouse position
        if (mx, my) != (px, py):
            # set the previous mouse position to the current mouse position
            px, py = (mx, my)

            # create a vector from the mouse location with a random offset on the x and y axis
            rx, ry = (mx + random.randint(-40, 40), my + random.randint(-40, 40))

            # create a point using pygame's vector2
            p = pygame.math.Vector2(rx, ry)

            # insert the point into the quadtree
            QT.insert(p)

    # clear the primary display by filling it with black
    PD.fill((0, 0, 0))
    QT.show()

    # if right mouse button is pressed then query quadtree
    if mb[2] and not mb[0]:
        # the query will take the form of a pygame.Rect (x, y, w, h)
        queryRect = pygame.Rect(mx - 100, my - 100, 200, 200)
        points = [] # this list will contain the points returned by the query

        # query the quadtree :D
        QT.query(queryRect, points)

        # draw query rectangle
        pygame.draw.rect(PD, (0, 255, 0), queryRect, 1)

        # highlight the points by over writing QT.show() points with queried points
        for p in points:
            pygame.draw.circle(PD, (255, 0, 255), (int(p.x), int(p.y)), 3)

    pygame.display.update()