基于前缀树(Trie)实现Python目录结构生成的性能优化问询

首先，我得说你的原生实现思路是对的，但性能瓶颈在于每次计算前缀频率都要遍历所有文件名，这在文件数量多的时候会非常慢。先给你一个优化后的原生实现，不用第三方库就能大幅提升性能，然后再讲如何用datrie或pytrie这类高效库来实现。

优化后的原生实现（性能大幅提升） #

核心是先一次性统计所有前缀的出现次数，而不是每次计算都遍历所有文件。这样时间复杂度从O(n²)降到了O(所有文件名的总长度)，效果立竿见影：

from collections import defaultdict

def count_prefixes(words):
    """一次性统计所有可能前缀的出现次数"""
    prefix_counts = defaultdict(int)
    for word in words:
        for i in range(1, len(word) + 1):
            prefix = word[:i]
            prefix_counts[prefix] += 1
    return prefix_counts

def make_trie_optimized(words, max_freq):
    prefix_counts = count_prefixes(words)
    root = dict()
    for word in words:
        current_dict = root
        for i in range(len(word)):
            letter = word[i]
            current_prefix = word[:i+1]
            if prefix_counts[current_prefix] > max_freq:
                # 前缀对应的文件数超过阈值，继续拆分目录
                current_dict = current_dict.setdefault(letter, {})
            else:
                # 达到阈值，标记为终止节点，停止处理当前文件
                current_dict[letter] = "end"
                break
    return root

# 原路径查询函数可以直接复用
def get_path(image_id, trie):
    result = ""
    current_dict = trie
    for i in range(len(image_id)):
        letter = image_id[i]
        if letter in current_dict:
            result += letter + "/"
            if current_dict[letter] == "end":
                break
            current_dict = current_dict[letter]
    return result

# 测试示例
filenames = ["111", "112", "1341", "2213", "2131", "22222", "11111"]
trie = make_trie_optimized(filenames, 2)
print("优化后的Trie结构:", trie)
for fname in filenames:
    print(f"{fname} -> {get_path(fname, trie)}")

用datrie实现高效Trie #

datrie是基于双数组的Trie实现，内存占用小、查询速度极快。它需要先定义字符集（这里是数字），我们可以用它存储每个前缀的状态（是否需要继续拆分）：

import datrie
from collections import defaultdict

def build_datrie(words, max_freq):
    # 先统计所有前缀的出现次数
    prefix_counts = defaultdict(int)
    for word in words:
        for i in range(1, len(word)+1):
            prefix = word[:i]
            prefix_counts[prefix] += 1
    
    # 定义字符集：0-9数字
    charset = '0123456789'
    # 创建Trie，值为bool：True=继续拆分，False=终止
    trie = datrie.Trie(charset)
    
    for word in words:
        current_prefix = ""
        for c in word:
            current_prefix += c
            count = prefix_counts[current_prefix]
            if count > max_freq:
                # 需要继续拆分，标记为True（如果还没存的话）
                if current_prefix not in trie:
                    trie[current_prefix] = True
            else:
                # 达到阈值，标记为终止，停止处理当前文件
                if current_prefix not in trie:
                    trie[current_prefix] = False
                break
        # 处理那些整个文件名都需要作为终止节点的情况
        if len(current_prefix) == len(word) and current_prefix not in trie:
            trie[current_prefix] = False
    return trie

def get_path_datrie(image_id, trie):
    path = ""
    current_prefix = ""
    for c in image_id:
        current_prefix += c
        path += c + "/"
        # 遇到终止节点就停止
        if current_prefix in trie and not trie[current_prefix]:
            break
    return path

# 测试
filenames = ["111", "112", "1341", "2213", "2131", "22222", "11111"]
trie = build_datrie(filenames, 2)
for fname in filenames:
    print(f"{fname} -> {get_path_datrie(fname, trie)}")

用pytrie实现 #

pytrie的SortedPrefixTrie支持便捷的前缀查询，我们可以直接把前缀的计数存在Trie里，查询时逐步判断：

from pytrie import SortedPrefixTrie
from collections import defaultdict

def build_pytrie(words, max_freq):
    prefix_counts = defaultdict(int)
    for word in words:
        for i in range(1, len(word)+1):
            prefix = word[:i]
            prefix_counts[prefix] += 1
    
    # 创建PrefixTrie，存储每个前缀的计数
    trie = SortedPrefixTrie()
    for prefix, count in prefix_counts.items():
        trie[prefix] = count
    return trie

def get_path_pytrie(image_id, trie, max_freq):
    path = ""
    current_prefix = ""
    for c in image_id:
        current_prefix += c
        count = trie[current_prefix]
        path += c + "/"
        if count <= max_freq:
            break
    return path

# 测试
filenames = ["111", "112", "1341", "2213", "2131", "22222", "11111"]
trie = build_pytrie(filenames, 2)
for fname in filenames:
    print(f"{fname} -> {get_path_pytrie(fname, trie, 2)}")

关键逻辑说明 #

不管用哪种实现，核心逻辑都是一致的：

• 统计所有前缀的出现次数：即有多少个文件名以该前缀开头
• 构建Trie时的终止判断：如果当前前缀对应的文件数≤阈值（这里是2），就标记为终止节点，不再继续拆分目录
• 路径查询：遍历文件名的每个字符，直到遇到终止节点，拼接出目录路径

这样就能保证每个最终目录下的文件数不超过2，同时性能比原生实现提升很多。

内容的提问来源于stack exchange，提问作者user1981275