python文件搜索脚本

因为Windows自带搜索功能实在太慢，尝试用python写了文件搜索的脚本。

主要使用了：

• os.walk(root_dir) 用来获得根目录下全部子目录与文件
• list.sort() 列表排序，是py自带排序算法，比自己写的冒泡快得多
• N-gram字符串比对算法，代码是自己写的不过思路就是N-gram的思路，把字符串切成等长的片段去做比较

"""
尝试使用N-gram算法对全盘文件/文件夹进行搜索
"""

import os


def calculate_cuts(str_in: str, n: int) -> int:
    # 计算切割次数
    length = len(str_in)
    if length <= n:
        return 1
    else:
        return length - n + 1


def cut(str_in: str, n: int) -> list:
    # 切割，返回得到的片段
    cuts = calculate_cuts(str_in, n)
    if cuts == 1:
        return [str_in]
    else:
        result_lst = list()
        for i in range(cuts):
            result_lst.append(str_in[i:i + n:1])
        return result_lst


def similarity_calculator(str_main: str, str_sub: str, n: int) -> float:
    # 计算两个字符串的相似度

    if str_main == str_sub:
        return 1.0
    if str_main == '' or str_sub == '':
        return 0.0

    # 切割
    main_lst = cut(str_main, n)
    sub_lst = cut(str_sub, n)

    # 计算片段种类总数与共享片段数
    sum_lst = list()
    for i in main_lst:
        sum_lst.append(['main', i])
    for i in sub_lst:
        if ['main', i] not in sum_lst:
            sum_lst.append(['sub', i])
    sum_counter = len(sum_lst)

    shared_counter = 0
    for i in main_lst:
        if i in sub_lst:
            shared_counter += 1

    return shared_counter / sum_counter


root_dir_input = input("在哪个根目录里搜索？输入C、D或E\n>>>")
root_dir = root_dir_input.capitalize() + ":\\.."

print("\n\n开始读取文件...\n\n")

files_lst = list()
for dirpath, dirnames, filenames in os.walk(root_dir):
    if len(files_lst) % 1000 == 0:
        print(f"读取进度: {len(files_lst)}")  # 报告进度
    # print(f"读取进度: {len(files_lst)}")
    for dirname in dirnames:
        files_lst.append([f"子目录: {os.path.join(dirpath, dirname)}", dirname])
    for filename in filenames:
        files_lst.append([f"文件: {os.path.join(dirpath, filename)}", filename])

print(f"\n\n读取完成，共{len(files_lst)}个文件/目录。\n\n")

str_input = input("请输入字符串以查询\n>>>")
n_str = input("请输入粒度N，默认为输入字符串长度，直接按Enter以采用默认值\n>>>")
n = int(n_str) if n_str != '' else len(str_input)

print("\n开始计算相似度...\n")

result_lst = list()
for i in range(len(files_lst)):
    [full_dir, name] = files_lst[i]
    if i % (len(files_lst) // 20) == 0:
        print(f"计算进度: {i // (len(files_lst) // 100)}%")
    similarity = similarity_calculator(str_input, name, n)
    if similarity > 0:  # 如果相似度为0就不需要加进最终列表了
        result_lst.append([full_dir, name, similarity])

print("\n相似度计算完成，进行排序...\n")

result_lst.sort(key=lambda x: x[2], reverse=True)  # 降序

print("排序完成，搜索结果如下：")
for i in range(len(result_lst)):
    similarity = result_lst[i][2]
    if similarity == 0:
        break
    print(f"{i + 1}\t{result_lst[i][1]} (similarity: {format(similarity, '.2f')}), {result_lst[i][0]}\n")

测试效果：比Windows自带搜索要快，主要耗时间在第一步读取根目录下文件与子目录