from numpy import *

def loadDataSet():
    return [[1, 3, 4], [2, 3, 5], [1, 2, 3, 5], [2, 5]]

dataSet=loadDataSet()
dataSet

def createC1(dataSet): #生成只包含一个元素的项集的集合C1
    C1 = []
    for transaction in dataSet:
        for item in transaction:
            if not [item] in C1:
                C1.append([item])
                
    C1.sort()
    return map(frozenset, C1)#use frozen set so we
                            #can use it as a key in a dict    
C1=createC1(dataSet)
C1
D=map(set,dataSet)
D

#计算Ck的项集的在数据集D中的支持度support,并检查数据确保每个项集都是频繁的
def scanD(D, Ck, minSupport): 
    """
    D:Lst[frozenset]
    Ck:Lst[frozenset]
    minSupport:int
    """
    ssCnt = {} #统计Ck中的项在D中的项中出现的次数
    for tid in D:
        for can in Ck:
            if can.issubset(tid): #如果can是tid的子集
                ssCnt[can]=ssCnt.get(can,0)+1
#            if can.issubset(tid):
#                if not ssCnt.has_key(can): ssCnt[can]=1
#                else: ssCnt[can] += 1
    numItems = float(len(D)) #D中共有多少条记录
    retList = []
    supportData = {}
    for key in ssCnt:
        support = ssCnt[key]/numItems #计算支持度：数据集中包含该项集的记录所占的比例
        if support >= minSupport:
            retList.append(key)
#            retList.insert(0,key)
        supportData[key] = support
    return retList, supportData

L1,suppData0=scanD(D,C1,0.5)
print(L1)
print(suppData0)

def aprioriGen(Lk, k): #creates Ck：这里k表示一个项集中包含几个元素
    """
    Lk: List(frozenset)
    """
    retList = []
    lenLk = len(Lk)
    for i in range(lenLk):
        for j in range(i+1, lenLk):  
            L1 = list(Lk[i])[:k-2]; L2 = list(Lk[j])[:k-2]
            L1.sort(); L2.sort()
            if L1==L2: #if first k-2 elements are equal
                retList.append(Lk[i] | Lk[j]) #set union
                """
                判断Lk集合的前k-2个元素相同，然后将Lk[i],Lk[j]合并
                这就能保证后面生成的L(k+1)集合中的项集中的元素个数为k+1
                
                这里k-2用于定位，举个栗子，当k=3，Lk=[{1,2},{1,3},{1,5},{2,3},{2,5},{3,5}]
                此时就判断Lk[i][1]和Lk[][1]即项集的第1个元素是否相等来决定是否将它们合并
                从而retList=[{1,2,3},{1,3,5},{1,2,5},{2,3,5}]
                
                若k=4，就判断Lk[i]和Lk[j]的前2个元素是否相等，
                Lk是输入，所有项集的元素都是k-1个
                因为我们要生成含有k个元素的所有项集的集合retList
                所以需要前k-2个元素相等，最后一个元素不相等，这样合并起来刚好有k个元素
                """
    return retList



def apriori(dataSet, minSupport = 0.5):
    C1 = createC1(dataSet)
    D = map(set, dataSet)
    L1, supportData = scanD(D, C1, minSupport)
    L = [L1]
    k = 2
    #下面利用L(k-1)来构建L(k),并添加新支持度字典数据
    while (len(L[k-2]) > 0): #当集合中项的个数大于0时
        Ck = aprioriGen(L[k-2], k) #构建一个k个项组成的候选项集的列表
        Lk, supK = scanD(D, Ck, minSupport)#scan DB to get Lk 
        supportData.update(supK) #将生成的新的支持度的字典添加进原来的包含支持度的字典中
        L.append(Lk) #
        k += 1 #k用于定位到相应的L的位置上，即当计算L2时定位到L[0],因为L[0]表示的是L1
    return L, supportData

L,suppData=apriori(dataSet,minSupport = 0.5)
print(L)
print(suppData)

aprioriGen(L[0], 2)
L,suppData=apriori(dataSet,minSupport=0.7)
print(L)
print(suppData)


def generateRules(L, supportData, minConf=0.7):  #supportData is a dict coming from scanD
    bigRuleList = []
    for i in range(1, len(L)):#only get the sets with two or more items
        i=2
        for freqSet in L[i]:
            H = [frozenset([item]) for item in freqSet]
            if (i > 1):
                rulesFromConseq(freqSet, H, supportData, bigRuleList, minConf)
                """
                当项集中的元素个数大于2时，调用rulesFromConseq
                rulesFromConseq：
                调用aprioriGen来生成conseq，便于生成不重复的关联规则
                根据关联规则调用calcConf计算conf
                """
            else:
                calcConf(freqSet, H, supportData, bigRuleList, minConf)
                """
                当项集中的元素个数等于2时，调用calcConf(筛选大于minconf的项集)
                """
    return bigRuleList         

def calcConf(freqSet, H, supportData, brl, minConf=0.7):
    prunedH = [] #create new list to return
    for conseq in H:
        conf = supportData[freqSet]/supportData[freqSet-conseq] #calc confidence
        if conf >= minConf: 
            print(freqSet-conseq,‘-->‘,conseq,‘conf:‘,conf)
            brl.append((freqSet-conseq, conseq, conf))
            prunedH.append(conseq)
    return prunedH

def rulesFromConseq(freqSet, H, supportData, brl, minConf=0.7):
    m = len(H[0])
    if (len(freqSet) > (m + 1)): #try further merging
        Hmp1 = aprioriGen(H, m+1)#create Hm+1 new candidates
        Hmp1 = calcConf(freqSet, Hmp1, supportData, brl, minConf)
        if (len(Hmp1) > 1):    #need at least two sets to merge
            #再对大于1个元素的频繁项集再细分计算关联规则
            rulesFromConseq(freqSet, Hmp1, supportData, brl, minConf)
            
def pntRules(ruleList, itemMeaning):
    for ruleTup in ruleList:
        for item in ruleTup[0]:
            print itemMeaning[item]
        print "           -------->"
        for item in ruleTup[1]:
            print itemMeaning[item]
        print "confidence: %f" % ruleTup[2]
        print       #print a blank line

L,suppData=apriori(dataSet,minSupport = 0.5)
rules=generateRules(L,suppData,minConf=0.7)

mushDatSet=[line.split() for line in open(‘mushroom.dat‘).readlines()]
L,suppData=apriori.apriori(mushDatSet,minSupport=0.3)
for item in L[1]:
    if item.intersection(‘2‘):print(item) #查看与特征值2相交的item