naivebayes.py

naivebayes.py

import csv
trainReader = csv.reader(open('traindata.csv', 'rb'))
nya = []
for row in trainReader:
    nya.append(row)
# train data
noarr = filter(lambda row: row[-1] == 'No', nya)
yesarr = filter(lambda row: row[-1] == 'Yes', nya)


testReader = csv.reader(open('testdata.csv', 'rb'))
testdata = []
for row in testReader:
    testdata.append(row)


# functions
def p_y(char,i,yesarr):
    return 1.0*len(filter(lambda row: row[i] == char, yesarr))/len(yesarr)

def pstr_y(arr,basearr):
    res = 1.0
    for i,val in enumerate(arr):
        res = res*p_y(val,i,basearr)
    return res

def pred(row):
    rr = row[0:-1]
    return pstr_y(rr,yesarr) >= pstr_y(rr,noarr)

# apply
for row in testdata:
    print 'Yes' if pred(row) else 'No'