python实现fp-tree算法的代码

您好，以下是Python实现FP-Growth算法的代码，供您参考：

# 定义FP树的节点
class TreeNode:
    def __init__(self, name, count, parent):
        self.name = name
        self.count = count
        self.parent = parent
        self.children = {}
        self.next = None
        
    def increase(self, count):
        self.count += count
        
    def display(self, ind=1):
        print('  '*ind, self.name, ' ', self.count)
        for child in self.children.values():
            child.display(ind+1)

# 构建FP树
def build_FP_tree(dataset, min_support):
    freq_items = {}
    for trans in dataset:
        for item in trans:
            freq_items[item] = freq_items.get(item, 0) + dataset[trans]
    freq_items = {k:v for k,v in freq_items.items() if v >= min_support}
    if len(freq_items) == 0:
        return None, None
    for item in freq_items:
        freq_items[item] = [freq_items[item], None]
    root = TreeNode('Root', 1, None)
    for trans, count in dataset.items():
        localD = {}
        for item in trans:
            if item in freq_items:
                localD[item] = freq_items[item][0]
        if len(localD) > 0:
            ordered_items = [v[0] for v in sorted(localD.items(), key=lambda p: p[1], reverse=True)]
            update_FP_tree(ordered_items, root, freq_items, count)
    return root, freq_items

# 更新节点和链表
def update_FP_tree(items, node, freq_items, count):
    if items[0] in node.children:
        node.children[items[0]].increase(count)
    else:
        node.children[items[0]] = TreeNode(items[0], count, node)
        if freq_items[items[0]][1] == None:
            freq_items[items[0]][1] = node.children[items[0]]
        else:
            update_links(freq_items[items[0]][1], node.children[items[0]])
    if len(items) > 1:
        update_FP_tree(items[1:], node.children[items[0]], freq_items, count)
  
# 更新连接节点
def update_links(node, target_node):
    while (node.next != None):
        node = node.next
    node.next = target_node

# 生成频繁项集的条件模式基
def find_prefix_path(node):
    cond_pats = {}
    while (node != None):
        prefix = []
        ascend_FP_tree(node, prefix)
        if len(prefix) > 1:
            cond_pats[tuple(prefix[1:])] = node.count
        node = node.next
    return cond_pats

# 回溯FP树，生成前缀
def ascend_FP_tree(node, prefix):
    if node.parent != None:
        prefix.append(node.name)
        ascend_FP_tree(node.parent, prefix)

# 递归查找频繁项集
def mine_FP_tree(freq_items, header_table, min_support, prefix, frequentItemList):
    items = [v[0] for v in sorted(header_table.items(), key=lambda p:p[1])]
    for item in items:
        new_freq_set = prefix.copy()
        new_freq_set.add(item)
        support = header_table[item][0]
        frequentItemList.append((new_freq_set, support))
        cond_patt_bases = find_prefix_path(header_table[item][1])
        cond_tree, cond_items = build_FP_tree(cond_patt_bases, min_support)
        if cond_tree != None:
            mine_FP_tree(cond_items, cond_tree, min_support, new_freq_set, frequentItemList)

# 封装函数
def FP_growth(dataset, min_support):
    freq_items = {}
    for trans in dataset:
        for item in trans:
            freq_items[item] = freq_items.get(item, 0) + dataset[trans]
    freq_items = {k:v for k,v in freq_items.items() if v >= min_support}
    if len(freq_items) == 0:
        return None
    for k in freq_items:
        freq_items[k] = [freq_items[k], None]
    header_table = {}
    for k in freq_items:
        header_table[k] = [freq_items[k][0], None]
    for tran, count in dataset.items():
        localD = {}
        for item in tran:
            if item in freq_items:
                localD[item] = freq_items[item][0]
        if len(localD) > 0:
            ordered_items = [v[0] for v in sorted(localD.items(), key=lambda p: p[1], reverse=True)]
            update_FP_tree(ordered_items, root, freq_items, count)
    frequentItemList = []
    mine_FP_tree(freq_items, header_table, min_support, set(), frequentItemList)
    return frequentItemList

注意：本算法需要输入数据集和最小支持度，返回频繁项集列表