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()