【算法优化】记一次不太成功的文本相似性去重算法优化实践【续】

0111001111110101011100100000111100110110010011011000010101110110

0111001111110101
0111001000001111
0011011001001101
1000010101110110

import time
import json
import numpy as np
from typing import List

# 常量配置
sim_thr = 0.85 * 64    # 85%相似度阈值

# 加载1.7万的文章id及simhash值
with open('./article_simhash_17k.json') as f:
    data = json.load(f)
    
print(len(data))
start = time.time()
new_data = {}
for item in data:
    simhash, article_id = item['simhash'], item['id']
    simhash = ['0']*(64-len(simhash)) + list(simhash)
    # 不足64位的前面补0，文章id放到第一个值
    item = [article_id] + simhash
    simhash = ''.join(simhash)
    for i in range(4):            # 这个方式耗时：28s（会有一定的精度损失）
        new_data.setdefault(str(i)+simhash[i*16:i*16+16], []).append(item)
    # for i in range(8):            # 这个方式整体时间：39s
    #     if i*8+16 <= 64:
    #         new_data.setdefault(str(i)+simhash[i*8:i*8+16], []).append(item)

all_np_data = [np.array(vals) for vals in new_data.values() if len(vals) > 1]
print('time: ', time.time()-start, len(new_data), len(all_np_data))

def cluster(np_data) -> List[List[str]]:
    """将相似的文章聚类在一起"""
    results = []
    while np_data.shape[0] > 1:
        curr_row = np_data[0]
        np_data = np_data[1:]
        counts = np.count_nonzero(curr_row == np_data, axis=1)
        if len(counts) > 0:
            ls = list(np_data[counts > sim_thr][:, 0])
            ls.append(curr_row[0])
            results.append(ls)
        np_data = np_data[counts <= sim_thr]

    return results

all_cls = []   # List[List[str]], 文章id的初步聚类结果，每个类可以包含多个文章id
for np_data in all_np_data:
    all_cls += cluster(np_data)     # TODO 这里是可以并行的，待优化（不过这里耗时不是太多）

results = {item['id']: []  for item in data}   # 每个id对应一个类别，空列表表示没有分到具体的类别里
for cls_id, aids in enumerate(all_cls):    # 循环处理每个原始类别
    for aid in aids:      # 把该类的文章id都对应到类别id
        results[aid].append(cls_id)    # 一个文章可以包含多个初始分类

print('cluster time: ', time.time()-start, len(all_cls))

# 合并所有有交集的分类id
# TODO 这个步骤最耗时间，超过99%的时间都是消耗在这里
clses = [set(vals) for vals in results.values() if len(vals) > 1]
merge_cls = np.array(range(len(all_cls)))      # 每个类别对应的原始id
print('before merge, for len:', len(clses))
for i, cls in enumerate(clses):
    for cls_j in clses[i+1:]:
        # 集合计算比较慢
        # if len(cls.intersection(clses[j])) == 0:
        if len(cls.intersection(cls_j)) == 0:
            continue
        # 有交集则对应的全修改为最小值
        # cls = cls.union(cls_j)     # time: 39s
        cls.update(cls_j)          # time: 38s
    # 获取最小的类别
    merge_cls[list(cls)] = min(cls)

print('merge time', time.time()-start, len(set(merge_cls)))

# 计算每个分类id对应的文章id列表
cls_results = {}
for aid, _cls in results.items():
    if len(_cls) > 0:
        cls_id = merge_cls[_cls[0]]   # 合并后的id
        cls_results.setdefault(cls_id, []).append(aid)

# 生成最后的结果
# TODO 需要验证这里的相似文章的相似度怎么样
cls_results = [vals for vals in cls_results.values()]
print('hot max:', max([len(vals) for vals in cls_results]))

# 计算热度
cls_results = [(vals[0], len(vals)) for vals in cls_results]

# 补全热度为1的数据
for _id, _cls in results.items():
    if len(_cls) == 0:
        cls_results.append((_id, 1))

print('time: ', time.time() - start, ' len: ', len(results), len(cls_results))
assert len(data) == sum([val[1] for val in cls_results])