优化首尾数字相同的大区间计数代码 解决超时问题

问题背景 #

需要统计满足1 < n < m < 10^18的区间内，所有首尾数字相同的数的数量。原实现通过遍历区间每个数并转换为字符串判断首尾数字，耗时5秒，超出1秒时间限制，需优化。

原代码：

def Digits(n,m):
     c=0
     for i in range(n,m+1):
        i=str(i)
        if i[0]==i[-1]:
            c+=1
     return c

核心优化思路：用数学规律直接计算，避免遍历 #

遍历每个数的时间复杂度是O(m-n)，当区间跨度接近10^18时完全不可行。换用数学方法计算[1, x]内符合条件的数的数量，再通过count(m) - count(n)得到目标区间的结果。

步骤1：实现count(x)函数，计算[1, x]内符合条件的数 #

按数的位数分段计算：

• 1位数（1-9）：所有数都满足，共9个。
• k位数（k≥2）：
  • 首数字为d（1-9）时，尾数字也为d的数共有10^(k-2)个（中间k-2位可任意填0-9）。
  • 若x不是完整的k位数区间，单独计算该段内符合条件的数。

def count(x):
    if x < 1:
        return 0
    s = str(x)
    length = len(s)
    total = 9  # 1-9的数量
    # 处理所有长度小于当前length的多位数
    for k in range(2, length):
        total += 9 * 10 ** (k-2)
    # 处理当前length的数
    first_digit = int(s[0])
    # 首数字小于first_digit的情况：每个首数字对应10^(length-2)个符合条件的数
    total += (first_digit - 1) * 10 ** (length - 2)
    # 首数字等于first_digit的情况，计算≤x的符合条件的数
    min_len_num = 10 ** (length - 1)
    prefix = int(s[:-1])
    target_num = prefix * 10 + first_digit
    if target_num <= x:
        total += prefix - (min_len_num // 10) + 1
    else:
        total += prefix - (min_len_num // 10)
    return total

步骤2：计算目标区间的结果 #

因为原问题约束是1 < n < m，目标区间为[n+1, m]，直接用count(m) - count(n)得到结果。

完整优化代码 #

def count(x):
    if x < 1:
        return 0
    s = str(x)
    length = len(s)
    total = 9  # 1-9的数量
    # 处理所有长度小于当前length的多位数
    for k in range(2, length):
        total += 9 * 10 ** (k-2)
    # 处理当前length的数
    first_digit = int(s[0])
    total += (first_digit - 1) * 10 ** (length - 2)
    min_len_num = 10 ** (length - 1)
    prefix = int(s[:-1])
    target_num = prefix * 10 + first_digit
    if target_num <= x:
        total += prefix - (min_len_num // 10) + 1
    else:
        total += prefix - (min_len_num // 10)
    return total

def Digits(n, m):
    return count(m) - count(n)

复杂度说明 #

该方法的时间复杂度为O(log x)，仅需处理数的位数（最多18位），计算量极小，完全能在1秒内完成。

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