如何优化计算数组相似元素对的高时间复杂度代码？

嘿，我来帮你搞定这个代码优化的问题～ 先说说你原来的代码为啥效率低：它用了三重循环，时间复杂度是O(N³)，而且逻辑里重复计算了大量元素对，最后还要去重，这对于稍微大一点的N来说，运行速度会慢到离谱。

核心思路梳理 #

首先得搞清楚题目里的相似性规则：

• 两个元素相似如果数值差1
• 相似性有传递性 → 这意味着所有属于连续整数区间的元素，互相之间都是相似的。比如数组里有1、2、3，那1和3也是相似的，因为1和2相似，2和3相似，传递下来1和3就相似了。

所以我们不需要逐个检查元素对，而是可以用更聪明的数学方式计算：

1. 先统计每个数值在数组里出现的次数
2. 把连续的数值归为一组（比如{1,2,3}是一组，{5,6}是另一组）
3. 对于每个组，计算组内所有元素能组成的无序对数量：公式是total * (total - 1) // 2，其中total是组内元素的总个数（比如组里有3个元素，能组成3*2/2=3对）
4. 把所有组的结果加起来就是最终答案

优化后的代码实现 #

from collections import Counter

def SimilarElementsPairs(A, N):
    # 统计每个数字的出现次数，O(N)时间
    num_counts = Counter(A)
    # 把所有唯一数字排序，O(M log M)时间，M是不同数字的数量（M ≤ N）
    sorted_nums = sorted(num_counts.keys())
    total_pairs = 0
    current_group_total = 0
    
    for idx, num in enumerate(sorted_nums):
        # 如果当前数字和前一个连续，就加入当前组
        if idx > 0 and num == sorted_nums[idx-1] + 1:
            current_group_total += num_counts[num]
        else:
            # 先结算前一个组的对数
            if current_group_total > 0:
                total_pairs += current_group_total * (current_group_total - 1) // 2
            # 开启新组
            current_group_total = num_counts[num]
    # 别忘了结算最后一个组
    if current_group_total > 0:
        total_pairs += current_group_total * (current_group_total - 1) // 2
    
    return total_pairs

# 输入处理（适配Python3，原代码是Python2的写法，这里调整了）
N = int(input())
A = list(map(int, input().split()))
out_ = SimilarElementsPairs(A, N)
print(out_)

时间复杂度对比 #

• 原代码：O(N³)，N稍微大一点（比如N=1000）就完全跑不动
• 优化后代码：O(N + M log M)，其中M是数组中不同数字的数量，实际运行效率提升几个数量级都不止

举个例子验证：比如数组是[1,2,2,3]，连续组是{1,2,3}，总元素数是4，能组成4*3/2=6对，分别是(0,1),(0,2),(0,3),(1,2),(1,3),(2,3)，和题目要求的结果一致。

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