Python: naive bayes

#!/usr/bin/env python
#encoding: utf-8
'''
docstring for naive bayes
x: refer to attr
y: refer to cls
p(y|x) = p(x|y) * p(y) / p(x)
but here needn't clac p(x)
'''

from __future__ import division

def calc_prob_cls(train, cls_val, cls_name='class'):
    '''
    calculate the prob. of class: cls
    '''
    cnt = 0
    for e in train:
        if e[cls_name] == cls_val:
            cnt += 1

    return cnt / len(train)

def calc_prob(train, cls_val, attr_name, attr_val, cls_name='class'):
    '''
    calculate the prob(attr|cls)
    '''
    cnt_cls, cnt_attr = 0, 0
    for e in train:
        if e[cls_name] == cls_val:
            cnt_cls += 1
            if e[attr_name] == attr_val:
                cnt_attr += 1

    return cnt_attr / cnt_cls

def calc_NB(train, test, cls_y, cls_n):
    '''
    calculate the naive bayes
    '''
    prob_y = calc_prob_cls(train, cls_y)
    prob_n = calc_prob_cls(train, cls_n)
    for key, val in test.items():
        print '%10s: %s' % (key, val)
        prob_y *= calc_prob(train, cls_y, key, val)
        prob_n *= calc_prob(train, cls_n, key, val)

    return {cls_y: prob_y, cls_n: prob_n}

if __name__ == '__main__':
    #train data
    train = [
        {"outlook":"sunny", "temp":"hot", "humidity":"high", "wind":"weak", "class":"no" },
        {"outlook":"sunny", "temp":"hot", "humidity":"high", "wind":"strong", "class":"no" },
        {"outlook":"overcast", "temp":"hot", "humidity":"high", "wind":"weak", "class":"yes" },
        {"outlook":"rain", "temp":"mild", "humidity":"high", "wind":"weak", "class":"yes" },
        {"outlook":"rain", "temp":"cool", "humidity":"normal", "wind":"weak", "class":"yes" },
        {"outlook":"rain", "temp":"cool", "humidity":"normal", "wind":"strong", "class":"no" },
        {"outlook":"overcast", "temp":"cool", "humidity":"normal", "wind":"strong", "class":"yes" },
        {"outlook":"sunny", "temp":"mild", "humidity":"high", "wind":"weak", "class":"no" },
        {"outlook":"sunny", "temp":"cool", "humidity":"normal", "wind":"weak", "class":"yes" },
        {"outlook":"rain", "temp":"mild", "humidity":"normal", "wind":"weak", "class":"yes" },
        {"outlook":"sunny", "temp":"mild", "humidity":"normal", "wind":"strong", "class":"yes" },
        {"outlook":"overcast", "temp":"mild", "humidity":"high", "wind":"strong", "class":"yes" },
        {"outlook":"overcast", "temp":"hot", "humidity":"normal", "wind":"weak", "class":"yes" },
        {"outlook":"rain", "temp":"mild", "humidity":"high", "wind":"strong", "class":"no" },
        ]   
    #test data
    test = {"outlook":"sunny","temp":"cool","humidity":"high","wind":"strong"}

    #calculate
    print calc_NB(train, test, 'yes', 'no')

参考：http://www.coder4.com/archives/1511