knn.py

import csv
import math
import operator

def euclideanDistance(p1, p2, length):
  distance = 0
  for x in range(length):
    distance += pow(p1[x] - p2[x], 2)
  return math.sqrt(distance)

def getNeighbors(trainingSet, datapoint, k):
  distances = []
  length = len(datapoint) - 1
  for x in range(len(trainingSet)):
    dist = euclideanDistance(trainingSet[x], datapoint, length)
    distances.append((trainingSet[x], dist))
  distances.sort(key=operator.itemgetter(1))
  neighbors = []
  for x in range(k):
    neighbors.append(distances[x][0])
  return neighbors

def tallyResults(neighbors):
  classVotes = {}
  for x in range(len(neighbors)):
    response = neighbors[x][-1]
    if response in classVotes:
      classVotes[response] += 1
    else:
      classVotes[response] = 1
  sortedVotes = sorted(classVotes.iteritems(), key = operator.itemgetter(1), reverse = True)
  return sortedVotes[0][0]

trainingSet = [[5.1, 3.5, 1.4, 0.2, 'Iris-setosa'], [4.9, 3.0, 1.4, 0.2, 'Iris-setosa'], [4.7, 3.2, 1.3, 0.2, 'Iris-setosa'], [7.0, 3.2, 4.7, 1.4, 'Iris-versicolor'], [6.4, 3.2, 4.5, 1.5, 'Iris-versicolor'], [6.9, 3.1, 4.9, 1.5, 'Iris-versicolor'], [6.3, 3.3, 6.0, 2.5, 'Iris-virginica'], [5.8, 2.7, 5.1, 1.9, 'Iris-virginica'], [7.1, 3.0, 5.9, 2.1, 'Iris-virginica']]

predictions = []
k = 3
datapoint = [5.7,2.8,4.1,1.3,'Iris-versicolor']
neighbors = getNeighbors(trainingSet, datapoint, k)
results = tallyResults(neighbors)
print results