FP-growth树

#coding=utf-8
from numpy import *
class tree_node:
    '''树节点'''
    def __init__(self,name_value,num_occur,parent_node):
        '''初始些变量'''
        self.name=name_value
        self.count=num_occur
        self.node_link=None
        self.parent=parent_node
        self.children={}
    def inc(self,num_occur):
        #元素计数
        self.count+=num_occur
    def disp(self,ind=1):
        #显示树
        print ' '*ind,self.name,' ' ,self.count
        #空格就是为了表示结构从属
        for child in self.children.values():
            child.disp(ind+1)
            
def create_tree(init_set,min_sup=1):
    '''创建树'''
    header_table={}
    #头表
    #内容就是各元素总计的出现次数
    for trans in init_set:
        for item in trans:
            header_table[item]=header_table.get(item,0)+init_set[trans]
            print header_table
            #建立元素表并计数
            #init_set[trans]恒等于1
            #之所以多加一步init_set是想起到筛重的作用
    for k in header_table.keys():
        #遍历头表的元素，如果有小于MIN_sup的就从字典中删掉
        if header_table[k]<min_sup:
            del(header_table[k])
    freq_item_set=set(header_table.keys())
    #这一步是想判断删除小于支持数后的表是否为空了
    #如果空了就结束
    if len(freq_item_set)==0:return None,None
    for k in header_table:
        #把表里每一个值都做成列表的形式
        header_table[k]=[header_table[k],None]
    ret_tree=tree_node('Null set',1,None)
    #创建第一个树节点
    for tran_set,count in init_set.items():
        #遍历初始化的集合看集合中的元素是否在上面处理过的集合里
        #如果在就载入字典locald，并赋值为头表的计数值
        locald={}
        for item in tran_set:
            if item in freq_item_set:
                locald[item]=header_table[item][0]
        if len(locald)>0:
            #如果locald有元素
            #将元素按计数值从大到小排序并更新树
            ordered_items=[v[0] for v in sorted(locald.items(),key=lambda p:p[1],reverse=True)]
            update_tree(ordered_items,ret_tree,header_table,count)
    return ret_tree,header_table
def update_tree(items,tree,header_table,count):
    '''更新树'''
    #判断第一个元素是否作为节点存在
    #若存在则增加计数，不存在就建立新的节点同时更新头表
    if items[0] in tree.children:
        tree.children[items[0]].inc(count)
    else:
        tree.children[items[0]]=tree_node(items[0],count,tree)
        if header_table[items[0]][1]==None:
            header_table[items[0]][1]=tree.children[items[0]]
        else:
            update_header(header_table[items[0]][1],tree.children[items[0]])
    if len(items)>1:
        update_tree(items[1::],tree.children[items[0]],header_table,count)
def update_header(node_test,target_node):
    while (node_test.node_link != None):
        node_test=node_test.node_link
    node_test.node_link=target_node
def load_data():
    data=[['r','z','h','j','p'],
          ['z','y','x','w','v','u','t','s'],
          ['z'],
          ['r','x','n','o','s'],
          ['y','r','x','z','q','t','p'],
          ['y','z','x','e','q','s','t','m']]
    return data
def create_init_set(data):
    ret_dict={}
    for trans in data:
        ret_dict[frozenset(trans)]=1
    return ret_dict
def test():
    data=load_data()
    init_set=create_init_set(data)
    print init_set
    tree,header=create_tree(init_set,3)
    tree.disp()
    print header
test()