KNN algorithm in python

$knn.py

'''
	Best KNN algorithm implementation world has ever seen
'''
import math

def EuclideanDistance(a, b): # N dimention Euclidean distance implementation
	s = 0
	for i in range(len(a)):
		s += (a[i] - b[i]) ** 2
	# return math.sqrt(s)
	# no need for sqrt
	return s

class KNN:
	def __init__(self, k, training_data):
		self.k = k
		self.samples = training_data
	
	def predict(self, params):
		distances = []
		for sample in self.samples:
			distances.append((sample[1], EuclideanDistance(params, sample[0])))
		distances.sort(key = lambda a: a[1])
		distances = list(map(lambda a: a[0], distances[0:self.k]))
		return max(set(distances), key = distances.count)

Iris_setosa     = 0
Iris_versicolor = 1
Iris_virginica  = 2

def ClassNameToEnum(class_name):
	if class_name == "Iris-setosa":
		return Iris_setosa
	if class_name == "Iris-versicolor":
		return Iris_versicolor
	if class_name == "Iris-virginica":
		return Iris_virginica
	print("------------ Error: Unknown Class Name!! ------------")

def EnumToClassName(enum):
	if enum == Iris_setosa:
		return "Iris-setosa"
	if enum == Iris_versicolor:
		return "Iris-versicolor"
	if enum == Iris_virginica:
		return "Iris-virginica"
	print("------------ Error: Unknown Class Name!! ------------")

def ReadData():
	training_data = []
	training_data_file = open("iris_train.txt")

	is_first_line = True
	for line in training_data_file.readlines():
		if is_first_line:
			is_first_line = False
			continue
		items = line.replace('\n', '').split(',')
		training_data.append([list(map(float, items[0:-1])), ClassNameToEnum(items[-1])])
		
	training_data_file.close()

	test_data = []
	test_data_file = open("iris_test.txt")

	is_first_line = True
	for line in test_data_file.readlines():
		if is_first_line:
			is_first_line = False
			continue
		items = line.replace('\n', '').split(',')
		test_data.append(list(map(float, items)))
		
	test_data_file.close()

	return (training_data, test_data)

def main():
	# read data
	# format: [[float(param), float(param), ...], int(class)]
	(training_data, test_data) = ReadData()

	# choose K
	# People on the net said: "use odd numbers for K (4K XD).
	# and using sqrt(n) where n is number of samples; works most of the time"
	k = math.sqrt(len(training_data))
	kRounded = round(k)
	if kRounded % 2 == 0:
		if kRounded > k:
			k = kRounded - 1
		else:
			k = kRounded + 1
	else:
		k = kRounded

	knn = KNN(k, training_data)

	for test in test_data:
		print(EnumToClassName(knn.predict(test)))

if __name__ == "__main__":
	main()

iris_test.txt

sepal length,sepal width,petal length,petal width
5,3.5,1.3,0.3
4.5,2.3,1.3,0.3
4.4,3.2,1.3,0.2
5,3.5,1.6,0.6
5.1,3.8,1.9,0.4
4.8,3,1.4,0.3
5.1,3.8,1.6,0.2
4.6,3.2,1.4,0.2
5.3,3.7,1.5,0.2
5,3.3,1.4,0.2
5.5,2.6,4.4,1.2
6.1,3,4.6,1.4
5.8,2.6,4,1.2
5,2.3,3.3,1
5.6,2.7,4.2,1.3
5.7,3,4.2,1.2
5.7,2.9,4.2,1.3
6.2,2.9,4.3,1.3
5.1,2.5,3,1.1
5.7,2.8,4.1,1.3
6.7,3.1,5.6,2.4
6.9,3.1,5.1,2.3
5.8,2.7,5.1,1.9
6.8,3.2,5.9,2.3
6.7,3.3,5.7,2.5
6.7,3,5.2,2.3
6.3,2.5,5,1.9
6.5,3,5.2,2
6.2,3.4,5.4,2.3
5.9,3,5.1,1.8
  

iris_train.txt

sepal length,sepal width,petal length,petal width,class
5.1,3.5,1.4,0.2,Iris-setosa
4.9,3,1.4,0.2,Iris-setosa
4.7,3.2,1.3,0.2,Iris-setosa
4.6,3.1,1.5,0.2,Iris-setosa
5,3.6,1.4,0.2,Iris-setosa
5.4,3.9,1.7,0.4,Iris-setosa
4.6,3.4,1.4,0.3,Iris-setosa
5,3.4,1.5,0.2,Iris-setosa
4.4,2.9,1.4,0.2,Iris-setosa
4.9,3.1,1.5,0.1,Iris-setosa
5.4,3.7,1.5,0.2,Iris-setosa
4.8,3.4,1.6,0.2,Iris-setosa
4.8,3,1.4,0.1,Iris-setosa
4.3,3,1.1,0.1,Iris-setosa
5.8,4,1.2,0.2,Iris-setosa
5.7,4.4,1.5,0.4,Iris-setosa
5.4,3.9,1.3,0.4,Iris-setosa
5.1,3.5,1.4,0.3,Iris-setosa
5.7,3.8,1.7,0.3,Iris-setosa
5.1,3.8,1.5,0.3,Iris-setosa
5.4,3.4,1.7,0.2,Iris-setosa
5.1,3.7,1.5,0.4,Iris-setosa
4.6,3.6,1,0.2,Iris-setosa
5.1,3.3,1.7,0.5,Iris-setosa
4.8,3.4,1.9,0.2,Iris-setosa
5,3,1.6,0.2,Iris-setosa
5,3.4,1.6,0.4,Iris-setosa
5.2,3.5,1.5,0.2,Iris-setosa
5.2,3.4,1.4,0.2,Iris-setosa
4.7,3.2,1.6,0.2,Iris-setosa
4.8,3.1,1.6,0.2,Iris-setosa
5.4,3.4,1.5,0.4,Iris-setosa
5.2,4.1,1.5,0.1,Iris-setosa
5.5,4.2,1.4,0.2,Iris-setosa
4.9,3.1,1.5,0.1,Iris-setosa
5,3.2,1.2,0.2,Iris-setosa
5.5,3.5,1.3,0.2,Iris-setosa
4.9,3.1,1.5,0.1,Iris-setosa
4.4,3,1.3,0.2,Iris-setosa
5.1,3.4,1.5,0.2,Iris-setosa
7,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
5.5,2.3,4,1.3,Iris-versicolor
6.5,2.8,4.6,1.5,Iris-versicolor
5.7,2.8,4.5,1.3,Iris-versicolor
6.3,3.3,4.7,1.6,Iris-versicolor
4.9,2.4,3.3,1,Iris-versicolor
6.6,2.9,4.6,1.3,Iris-versicolor
5.2,2.7,3.9,1.4,Iris-versicolor
5,2,3.5,1,Iris-versicolor
5.9,3,4.2,1.5,Iris-versicolor
6,2.2,4,1,Iris-versicolor
6.1,2.9,4.7,1.4,Iris-versicolor
5.6,2.9,3.6,1.3,Iris-versicolor
6.7,3.1,4.4,1.4,Iris-versicolor
5.6,3,4.5,1.5,Iris-versicolor
5.8,2.7,4.1,1,Iris-versicolor
6.2,2.2,4.5,1.5,Iris-versicolor
5.6,2.5,3.9,1.1,Iris-versicolor
5.9,3.2,4.8,1.8,Iris-versicolor
6.1,2.8,4,1.3,Iris-versicolor
6.3,2.5,4.9,1.5,Iris-versicolor
6.1,2.8,4.7,1.2,Iris-versicolor
6.4,2.9,4.3,1.3,Iris-versicolor
6.6,3,4.4,1.4,Iris-versicolor
6.8,2.8,4.8,1.4,Iris-versicolor
6.7,3,5,1.7,Iris-versicolor
6,2.9,4.5,1.5,Iris-versicolor
5.7,2.6,3.5,1,Iris-versicolor
5.5,2.4,3.8,1.1,Iris-versicolor
5.5,2.4,3.7,1,Iris-versicolor
5.8,2.7,3.9,1.2,Iris-versicolor
6,2.7,5.1,1.6,Iris-versicolor
5.4,3,4.5,1.5,Iris-versicolor
6,3.4,4.5,1.6,Iris-versicolor
6.7,3.1,4.7,1.5,Iris-versicolor
6.3,2.3,4.4,1.3,Iris-versicolor
5.6,3,4.1,1.3,Iris-versicolor
5.5,2.5,4,1.3,Iris-versicolor
6.3,3.3,6,2.5,Iris-virginica
5.8,2.7,5.1,1.9,Iris-virginica
7.1,3,5.9,2.1,Iris-virginica
6.3,2.9,5.6,1.8,Iris-virginica
6.5,3,5.8,2.2,Iris-virginica
7.6,3,6.6,2.1,Iris-virginica
4.9,2.5,4.5,1.7,Iris-virginica
7.3,2.9,6.3,1.8,Iris-virginica
6.7,2.5,5.8,1.8,Iris-virginica
7.2,3.6,6.1,2.5,Iris-virginica
6.5,3.2,5.1,2,Iris-virginica
6.4,2.7,5.3,1.9,Iris-virginica
6.8,3,5.5,2.1,Iris-virginica
5.7,2.5,5,2,Iris-virginica
5.8,2.8,5.1,2.4,Iris-virginica
6.4,3.2,5.3,2.3,Iris-virginica
6.5,3,5.5,1.8,Iris-virginica
7.7,3.8,6.7,2.2,Iris-virginica
7.7,2.6,6.9,2.3,Iris-virginica
6,2.2,5,1.5,Iris-virginica
6.9,3.2,5.7,2.3,Iris-virginica
5.6,2.8,4.9,2,Iris-virginica
7.7,2.8,6.7,2,Iris-virginica
6.3,2.7,4.9,1.8,Iris-virginica
6.7,3.3,5.7,2.1,Iris-virginica
7.2,3.2,6,1.8,Iris-virginica
6.2,2.8,4.8,1.8,Iris-virginica
6.1,3,4.9,1.8,Iris-virginica
6.4,2.8,5.6,2.1,Iris-virginica
7.2,3,5.8,1.6,Iris-virginica
7.4,2.8,6.1,1.9,Iris-virginica
7.9,3.8,6.4,2,Iris-virginica
6.4,2.8,5.6,2.2,Iris-virginica
6.3,2.8,5.1,1.5,Iris-virginica
6.1,2.6,5.6,1.4,Iris-virginica
7.7,3,6.1,2.3,Iris-virginica
6.3,3.4,5.6,2.4,Iris-virginica
6.4,3.1,5.5,1.8,Iris-virginica
6,3,4.8,1.8,Iris-virginica
6.9,3.1,5.4,2.1,Iris-virginica