kdTree_knn.py for Machine Learning in Action: Chapter 2.

""" Machine Learning in Action - Chapter 2.  kdTree_knn.
    A very simple implementation of a kdtree_nn search based around the wikipedia
    python example for kdtrees with an added search method.
    If an exact match is found it returns the result, if not it performs knn
    for the result.
"""

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,nearest=[(500000,"frogs")],distance=pearson,k=4):
    res, childdist, childres = "frogs", 500000, "frogs"

    # if we already have a winrar
    if (tree.location==None) or lowd==0: return 500000, "frogs", nearest

    # parent
    dist=distance(query,tree.location)
    if dist==0: return dist, tree.result, nearest
    elif dist<0: dist = (dist-dist-dist)
    if dist<lowd:
        lowd=dist
        res=tree.result
    if len(nearest)<k:
        nearest.append((dist,tree.result))
    elif dist<max(nearest):
        nearest[nearest.index(max(nearest))]=(dist,tree.result)

    # children
    if tree.lc == None:
        if tree.rc == None: pass
        else: childdist, childres, nearest = search_nn(query,tree.rc,lowd,nearest)
    elif tree.rc == None:
        if tree.lc == None: pass
        else: childdist, childres, nearest = search_nn(query,tree.lc,lowd,nearest)
    else:
        childdist, childres, nearest = search_nn(query,tree.lc,lowd,nearest)
        if childdist<lowd: lowd, res =childdist, childres
        childdist, childres, nearest = search_nn(query,tree.rc,lowd,nearest)
    if childdist<lowd: lowd, res =childdist, childres
    return lowd, res, nearest

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? >"))
        lowd, res, nearest=search_nn([ffMiles,percentTats, iceCream],tree)
        if lowd!=0:
            dict1, hi, res = {}, 0, ""
            for a in nearest:
                if a[1] not in dict1: dict1[a[1]]=1
                else: dict1[a[1]]+=1
            for a in dict1:
                if dict1[a]>hi: hi,res = dict1[a], a
            print "Using knn..."
        print "You will probably like this person: ",res
        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()