Search code examples
python-3.xmath

How do I optimize my code to run with n <= 100 and a,b <= 10^5?


#How do I optimize my code to run with n <= 100 and a,b <= 10^5?

Here's the question

There are N test cases, each test case enters 2 integers a and b, summing all digits from a to b Input:

  • The first line is N
  • For the next N lines, enter the numbers a, b

Output is N lines with N results of the problem

Sample:

Input:

2
5 12
12 60

Output:

41
378

Note: At the first test case, we have the sum = 5+6+7+8+9+1+0+1+1+1+2 = 41

Here is my code, but it got Time Limit Exceeded

def main():
    def tonga(a):
        tong = 0
        while a > 0:
            tong += a % 10
            a //= 10
        return tong
    def main1():
        a, b = map(int, input().split())
        tong = 0
        for z in range(a, b + 1):
            if z < 10:
                tong += z
            else:
                tong += tonga(z)
        print(tong)
    n = int(input())
    for _ in range(n):
        main1()
main()

Solution

  • Compute the running sum of digits from 0 to each number up to 10^5. For example:

    A = [0] * 100001
    for i in range(1, 100001):
        A[i] = A[i-1] + sum(int(d) for d in str(i))
    

    Now, for each a and b, you can compute the sum of the digit sums from a to b in O(1) time: it's simply A[b] - A[a-1].

    For a solution that works without precomputation (for example if the input range is much larger than 10^5), you can find the sum of the digits from all the numbers from 1 to N in O((logN)^2) arithmetic operations. Based on https://math.stackexchange.com/questions/817038/sum-of-digits-of-number-from-1-to-n

    def dss(a):
        T, M = 0, 1
        while a > 0:
            q, r = a // 10, a % 10
            T += M * (45 * q + (r+1)*r//2 + (r+1)*sum(map(int, str(q))))
            a, M = q - 1, M*10
        return T
    
    def dig_sum_range(a, b):
        return dss(b) - dss(a-1)
    
    print(dig_sum_range(1, 1000000000000))