携程机器学习笔试编程第二题朴素贝叶斯

# -*- coding: utf-8 -*-
"""
Spyder Editor

This is a temporary script file.
"""

def pretempe(train,test,indexlist):
    for index in indexlist:
        temp=[]
        for i in range(len(train)):
            temp.append(train[i][index])
        tmax,tmin=max(temp),min(temp)
        p1=tmin+(tmax-tmin)/3
        p2=tmax-(tmax-tmin)/3
        for i in range(len(train)):
            if temp[i]<p1: train[i][index]=0
            elif temp[i]<p2: train[i][index]=1
            else: train[i][index]=2
        if test[index]<p1: test[index]=0
        elif test[index]<p2: test[index]=1
        else: test[index]=2
    return train,test

if __name__=='__main__':
    n=9
    label=[0, 0, 0, 0, 1, 1, 1, 1, 1]
    train=[[0,0, 30, 450, 7],
[1, 1, 5,  500, 3],
[1 ,0, 10, 150, 1],
[0, 1, 40, 300, 6],
[1, 0, 20, 100, 10],
[0, 1, 25, 180, 12],
[0, 0, 32, 50 , 11],
[1, 0, 23, 120, 9],
[0, 0, 27, 200, 8]]
    test=[0, 0, 40, 180, 8]
    
    train,test=pretempe(train,test,[2,3,4])
    f_num=len(test)
    s_num=len(train)
    prior=[[0]*3 for _ in range(f_num)]
    for i in range(f_num):
        tmp=[]
        for k in range(s_num):
            tmp.append(train[k][i])
        for j in range(3):
            prior[i][j]=tmp.count(j)/s_num
    condition=[[[0]*3 for i in range(f_num)] for j in range(2)]
    label0=[]
    label1=[]
    for i in range(s_num):
        if label[i]==0: label0.append(i)
        else: label1.append(i)
    for i in range(f_num):
        p1=p2=p3=0
        for j in label0:
            
            if train[j][i]==0: p1+=1
            elif train[j][i]==1: p2+=1
            else: p3+=1
        condition[0][i][0]=p1/len(label0)
        condition[0][i][1]=p2/len(label0)
        condition[0][i][2]=p3/len(label0)
        p1=p2=p3=0
        for j in label1:
            if train[j][i]==0: p1+=1
            elif train[j][i]==1: p2+=1
            else: p3+=1
        condition[1][i][0]=p1/len(label1)
        condition[1][i][1]=p2/len(label1)
        condition[1][i][2]=p3/len(label1)
    l0,l1=len(label0)/s_num,len(label1)/s_num
    t1,t2,t3,t4,t5=test
    
    res1=l1*condition[1][0][t1]*condition[1][1][t2]*condition[1][2][t3]*condition[1][3][t4]*condition[1][4][t5]
    res1=res1/prior[0][t1]/prior[1][t2]/prior[2][t3]/prior[3][t4]/prior[4][t5]
    
    res0=l0*condition[0][0][t1]*condition[0][1][t2]*condition[0][2][t3]*condition[0][3][t4]*condition[0][4][t5]
    res0=res0/prior[0][t1]/prior[1][t2]/prior[2][t3]/prior[3][t4]/prior[4][t5]
    
    print(round(res1/res0,3))