Visualize K-D Tree on Mesh

'''
    WIP: 15_1207
    Script By: Dr. Gravitas

    Visualize K-D Tree Spacial Partitioning of selected Mesh.

'''
import pymel.core as pm
from maya import OpenMaya as om
import os
import fileinput
import pickle
import heapq
from collections import namedtuple
from operator import itemgetter
from pprint import pformat


def buildTreeAndHash( target ):
    for vertex in target.verts:
        vPoint = vertex.getPosition( space='world' )
        pos = ( vPoint.x, vPoint.y, vPoint.z )

        v2PDict[pos] = vertex

    tree = kdtree( v2PDict.keys() )
    return tree

class Node(namedtuple('Node', 'location tree_level axis left_child right_child')):
    def __repr__(self):
        return pformat(tuple(self))

def kdtree( point_list, depth=0 ):
    try:
        k = len( point_list[0] ) # assumes all points have the same dimension
    except IndexError as e: # if not point_list:
        return None
    # Select axis based on depth so that axis cycles through all valid values
    axis = depth % k

    # Sort point list and choose median as pivot element
    point_list = sorted( point_list, key=itemgetter(axis))
    median = len( point_list ) // 2 # choose median

    # Create node and construct subtrees
    return Node(
        location = point_list[median],
        tree_level = depth,
        axis = axis,
        left_child = kdtree(point_list[:median], depth + 1),
        right_child = kdtree(point_list[median + 1:], depth + 1)
    )

def buildVertSelectGroups( node ):
    currentGroup = []
    allGroups = []
    if node is not None:
        currentVert = v2PDict[node.location]
        currentGroup.extend( currentVert )
        leftGroup = buildVertSelectGroups( node.left_child )
        rightGroup = buildVertSelectGroups( node.right_child )

        currentGroup.extend(leftGroup[0])
        currentGroup.extend(rightGroup[0])
        allGroups.append(currentGroup)

        for vGroup in leftGroup[1]:
            allGroups.append( vGroup )
        for vGroup in rightGroup[1]:
            allGroups.append( vGroup )


    return currentGroup, allGroups


'''
    Main Body
'''
start_time = pm.date()
print '\n\n\n'
print '******************************************************************'
print '\t\t\t' + start_time
print '******************************************************************'
target = (pm.selected())[0].getShape()
pm.select( clear=True )
tree = ()
v2PDict = {}
tree = buildTreeAndHash( target )
selectionGroups = buildVertSelectGroups( tree )
selectionGroups[1].append( selectionGroups[0] )
selectionGroups = selectionGroups[1]

selectionGroups.sort(key=len, reverse=True)

for selectionGroup in selectionGroups:
    pm.select( selectionGroup )
    pm.refresh(f=True)
pm.select( clear=True )
pm.refresh(f=True)
tree = None
v2PDict = None
selectionGroups = None
end_time = pm.date()
print '******************************************************************'
print '\t\t\t' + start_time
print '\t\t\t' + end_time
print '******************************************************************'
#pm.error('Success!')