Search code examples
pythonalgorithmnumber-theory

Find b that (a+b) divisible to K


I have integer input: 0 < a, K, N < 10^9
I need to find all b numbers that satisfy:

  • a + b <= N
  • (a + b) % K = 0

For example: 10 6 40 -> [2, 8, 14, 20, 26]
I tried a simple brute force and failed (Time Limit Exceeded). Can anyone suggest answer? Thanks

a, K, N = [int(x) for x in input().split()]
count = 0
b = 1

while (a + b <= N):
    if ((a + b) % K) == 0:
        count+=1
        print(b, end=" ")
    b+=1

if (count == 0):
    print(-1)

Solution

  • The first condition is trivial in the sense that it just poses an upper limit on b. The second condition can be rephrased using the definition of % as

    a + b = P * K
    

    For some arbitrary integer P. From this, is simple to compute the smallest b by finding the smallest P that gives you a positive result for P * K - a. In other words

    P * K - a >= 0
    P * K >= a
    P >= a / K
    P = ceil(a / K)
    

    So you have

    b0 = ceil(a / K) * K - a
    b = range(b0, N + 1, K)
    

    range is a generator, so it won't compute the values up front. You can force that by doing list(b).

    At the same time, if you only need the count of elements, range objects will do the math on the limits and step size for you conveniently, all without computing the actual values, so you can just do len(b).