数据挖掘-FP_Growth算法

import itertools
from collections import defaultdict
import copy

#全局变量
res=[]

#从文件中加载数据
def load_data(path="data.txt"):
    res=[]
    with open(path,'rb') as f:
        for i in f:
            res.append(list(map(int,i.split())))
    print("读取文件并用list（）结构化数据：",res)
    return res

#生成有序、剔除非1-频繁项集中元素后的数据;    return (有序数据集，dict(1-频繁项集))
def generate_ordered_frequent_data(data,thre):
    data=copy.deepcopy(data)
    #print(data)
    d=defaultdict(int)
    for i in data:
        for j in i:
            d[j]+=1
    for i in range(len(data)):
        data[i].sort(key=lambda x:d[x],reverse=True)
        for j in range(len(data[i])):
            if d[data[i][j]]<thre:
                break
        if d[data[i][j]]<thre:
            data[i]=data[i][:j]
    #print("生成有序、剔除非1-频繁项集中元素后的数据集:",data)
    for i in list(d.keys()):#必须这么删除，否则 RuntimeError: dictionary changed size during iteration
        if d[i]<thre:
            del d[i]
    #print("生成频繁1-项集：",d)
    return data,d

#树节点定义
class Node:
    def __init__(self,val=-1):
        self.val=val
        self.cnt=1
        self.parent=None
        self.children={}
        self.next=None#header_tabel中用于连接下一个同名节点


#建树，传入有序数据集
def build_tree(data):
    root=Node()
    header_tabel= {}
    for i in data:
        p=root
        for j in i:
            if j in header_tabel:
                header_tabel[j][0]+=1
            else:
                header_tabel[j]=[1,None]
            if j in p.children:
                p=p.children[j]
                p.cnt+=1
            else:
                node=Node(j)
                node.parent=p
                p.children[j]=node
                header_tabel_node=header_tabel[j][1]
                if header_tabel_node==None:
                    header_tabel[j][1]=node
                else:
                    while header_tabel_node.next!=None:
                        header_tabel_node=header_tabel_node.next
                    header_tabel_node.next=node
                p=node
    l = (sorted(header_tabel, key=lambda x: header_tabel[x][0]))
    d={}
    for i in l:
        d[i]=header_tabel[i]
    header_tabel=d
    return root,header_tabel

#打印一下树康康
def print_root(root):
    q=[]
    q.append(root)
    print(root.val)
    while q:
        tmp=q.pop(0)
        for i in tmp.children:
            print(i,end=' ')
            q.append(tmp.children[i])
        print()

#判断传进来的root节点是否只有一条路径
def is_single_path(root):
    q=[]
    q.append(root)
    while q:
        tmp=q.pop(0)
        if len(tmp.children)>1:
            #print(tmp.children)
            return False
        for i in tmp.children:
            q.append(tmp.children[i])
    return True

#生成条件模式基
def generate_con_pattern_base(root,x,header_tabel):
    p=header_tabel[x][1]
    d={}
    while p!=None:
        cnt=p.cnt
        tmp=p.parent
        l=[]
        while tmp!=root:
            l.append(tmp.val)
            tmp=tmp.parent
        l.reverse()
        if l:
            #print(l)
            d[tuple(l)]=cnt
        p=p.next
    #print(x,d)
    return d

#生成小事务集
def generate_small_data_set(d):
    d=copy.deepcopy(d)
    l=[]
    for i in d:
        for j in range(d[i]):
            l.append(list(i))
    return l





#FP_Growth算法
def FP_Growth(a,data,thre):

    global res
    #生成有序数据集
    ordered_data,_=generate_ordered_frequent_data(data,thre)
    #print(a,data,"xxxx",ordered_data)

    #建树、header tabel
    root,header_tabel=build_tree(ordered_data)
    #print("",header_tabel)
    #print(header_tabel)

    #判断是否为单路径
    if is_single_path(root):#若为单路径，直接组合拼接，加入到res中
        #print("single path!")
        l=[]
        p=root
        while p!=None:
            if p.cnt>=thre and p.val!=-1:
                l.append(p.val)
            if list(p.children.keys()):
                p=p.children[list(p.children.keys())[0]]
            else:
                break
        for i in range(1,len(l)+1):
            tmp= list(itertools.combinations(l, i))#长度为i的l的不同组合
            for j in tmp:
                if list(j)+a not in res:
                    res.append(list(j)+a)
        if a not in res:
            res.append(a)
    else:#若不为单路径，准备递归
        for x in header_tabel:#此处header tabel已按从小到大排列
            i_con_pattern_base=generate_con_pattern_base(root,x,header_tabel)#得到条件模式基
            #print(i,i_con_pattern_base)
            if i_con_pattern_base:
                #print(x,"的子事务集如下：")
                small_data_set=generate_small_data_set(i_con_pattern_base)#根据条件模式基生成子数据集
                if a and a not in res:
                    res.append(a)
                FP_Growth(a+[x],small_data_set,thre)#递归咯！
            else:
                if a+[x] not in res:
                    res.append(a+[x])


    #print(root,header_tabel)
    #print_root(root)
    return

#保存到 result.txt文件中
def save_in_txt(minsup,path="result.txt"):

    with open(path,"a+") as f:
        f.write("\nminsup="+str(minsup))
        i=0
        n=len(res)
        k=1
        while i<n:
            j=i
            while i<n and len(res[i])==k:
                i+=1
            tmp=res[j:i]
            if tmp:
                f.write("\nL"+str(k)+":"+str(tmp))
                k+=1
            else:
                break

if __name__=="__main__":
    data_origin=load_data()
    minsup=0.2
    thre=len(data_origin)*minsup
    FP_Growth([],data_origin,thre)
    for i in range(len(res)):
        res[i].sort()
    res.sort(key=lambda x:(len(x),x[0]))
    print(res)
    save_in_txt(minsup)