kdTree_nn.py for Machine Learning in Action - Chapter 2

""" Machine Learning in Action - Chapter 2.  kdTree_nn.
    A very simple implementation of a kdtree_nn search based around the wikipedia
    python example for kdtrees with an added search method.
             -lolamontes69

    ToDo: Implement mixed with knn so that if zero distance found it returns
          that result, else it returns knn from a list of the k-lowest distances.
"""

from operator import itemgetter
from math import *

class Node:
    def __init__(self,location,result,left_child,right_child):
        self.location=location
        self.result=result
        self.lc=left_child
        self.rc=right_child


def kdtree(point_list, depth=0):
    try:
        k = len(point_list[0])
    except IndexError as e:
        return None
    axis = depth % k
    point_list.sort(key=itemgetter(axis))
    median = len(point_list) // 2                  # choose median
    return Node(location=point_list[median][0],
                result=point_list[median][1],
                left_child=kdtree(point_list[:median], depth + 1),
                right_child=kdtree(point_list[median + 1:], depth + 1))

def file2matrix(filename):
    returnMat=[]
    classLabelVector = []
    fr = open(filename)
    for line in fr.readlines():
        line = line.strip()
        listFromLine = line.split('\t')
        returnMat.append(listFromLine[0:3])
        classLabelVector.append(str(listFromLine[-1]))
    return returnMat, classLabelVector

def pearson(v1a, v2a):
    v1=[float(a) for a in v1a]
    v2=[float(a) for a in v2a]
    sum1 = sum(v1)
    sum2 = sum(v2)
    sum1Sq = sum([pow(v, 2) for v in v1])
    sum2Sq = sum([pow(v, 2) for v in v2])
    pSum = sum([v1[i]*v2[i] for i in range(len(v1))])
    num = pSum-(sum1*sum2/len(v1))
    den = sqrt((sum1Sq-pow(sum1, 2)/len(v1))*(sum2Sq-pow(sum2, 2)/len(v1)))
    if den == 0: return 0
    return 1.0-num/den

def search_nn(query,tree,lowd=500000,distance=pearson):
    res, childdist, childres = "frogs", 500000, "frogs"
    if (tree.location==None) or lowd==0: return 500000, "frogs"
    dist=distance(query,tree.location)
    if dist==0: return dist, tree.result
    if dist<0: dist = (dist-dist-dist)
    if dist<lowd:
        lowd=dist
        res=tree.result
    if tree.lc == None:
        if tree.rc == None: pass
        else: childdist, childres = search_nn(query,tree.rc,lowd)
    elif tree.rc == None:
        if tree.lc == None: pass
        else: childdist, childres = search_nn(query,tree.lc,lowd)
    else:
        childdist, childres = search_nn(query,tree.lc,lowd)
        if childdist<lowd: lowd, res =childdist, childres
        childdist, childres = search_nn(query,tree.rc,lowd)
    if childdist<lowd: lowd, res =childdist, childres
    return lowd, res

def main():
    datingDataMat, datingLabels = file2matrix('datingTestSet.txt')
    point_list=[]
    for a in range(len(datingDataMat)):
        point_list.append([list(datingDataMat[a]),datingLabels[a]])
    tree = kdtree(point_list)
    while 1:
        percentTats = float(raw_input("Percentage of time spent playing Video Games? >"))
        ffMiles = float(raw_input("Frequent flier miles earned per year? >"))
        iceCream = float(raw_input("Liters of ice cream consumed per year? >"))
        res=search_nn([ffMiles,percentTats, iceCream],tree)
        print "You will probably like this person: ",res[1]
        print "-"*44
        another=raw_input('Classify another?(y/n) >')
        if another=='n': break
        print "-"*44

def test():
    datingDataMat, datingLabels = file2matrix('datingTestSet.txt')
    point_list=[]
    for a in range(len(datingDataMat)):
        point_list.append([list(datingDataMat[a]),datingLabels[a]])
    tree = kdtree(point_list)
    # test each item in the dataset to see if it returns zero distance
    test_list=[]
    for a in point_list:
        print search_nn(a[0],tree),' ',

if __name__ == "__main__":
    main()