Revised: Counting subnumbers divisible by a given prime

Question

Assume we are given a number and a prime p. The goal is to count all subnumbers of number that are divisible by p. A subnumber of number is formed of any number of consecutive digits in the decimal representation of number.

Following answers to my first question about code to solve this problem, I improved my code as follows:

from collections import Counter

number = input()
p = int(input())

ans = 0

if p != 2 and p != 5:

    res = (int(number[i:]) % p for i in range(len(number)))
    C = Counter(res)

    for c in C:
        if c == 0:
            ans += (C[c] * (C[c] + 1)) // 2
        else:
            ans += ((C[c] - 1) * C[c]) // 2

elif p == 2:
    for i in range(len(number)):
        if number[i] in ('0', '2', '4', '6', '8'):
            ans += i + 1

else:
    for i in range(len(number)):
        if number[i] in ('0', '5'):
            ans += i + 1

print(ans)

Now, it passes more tests but yet there are several tests with time limit exceeded error.

Can this new solution be made faster?

Answer 1

If you say you get a 'time limit exceeded' message it would be interesting how large the string and the prime number and the time limit is. It seems that the prime is about \$10^9\$ , the string length is about \$10^6\$ and the time limit is about 2 seconds. You should use a profiler to find out where you spent the time in your program. I did this for a string of length \$10^4\$ . The program took about 11 seconds on my notebook and almost all of the time is spent for the calculation of res. The loop is executed \$10^4\$ times, so each calculation took about 1 ms. What is done during this millisecond? From the string of length \$10^4\$ a substring most of the time of length \$10^4\$ is taken and converted to an integer. This \$10^4\$ digit integer is divided by p to get the residue.

Can the performance of such a step be improved by using the results of a previous step?

Yes, I think so. The new substring differs from the old substring only by one character. So you have to convert only this one character to integer. The new residuum can be calculated now based on this number and the previously calculated values like the previous residue.

If you need the powers of \$10^i\$ in your calculation modulo p you should keep them reduced modulo p, otherwise they will become rather big (\$10^6\$ digit integers)

So the following is bad because you calculate with 1000000-digit numbers

p= some prime number of size (not length) 10**9
power=1
for _ in range(10**6):
    power=10*power
result=power % p

This one is much much better:

p= some prime number of size (not length) 10**9
power=1
for _ in range(10**6):
    power=(10*power)%p
result=power % p

because you calculate with 9-digit numbers

you can always measure the time your code needs with time.time

import time
starttime=time.time()

my code 

endtime=time.time()
print(endtime-starttime)

You will see the difference in the previous examples with the loops.

To testthe performance of the code one needs a 9-digit prime number and a \$10^6\$ digit string.

p=623456849
random.seed(12345678)
number=''.join(random.choices(string.digits, k=10**6))

The random.seed command guarantees that always the sam string is generated. The prime number p is from the internet. But for p one can use an arbitrary number that is not divisible by 2 or 5.

Answer 2

It is a bit weird to have two special cases and start with the case that it is neither of those. I would re-arrange your cases to

if p == 2:
    ...
elif p == 5:
    ...
else:
    ...

Unless you did this as a performance improvement, because a random prime is much more likely to be neither 2 nor 5, but in that case you should have some measurements to show that this is not a premature optimization and probably add a comment in the code.

---

If you iterate over something and need both the index and the value, use enumerate:

if p == 2:
    ans = sum(i for i, digit in enumerate(number, 1)
              if digit in {'0', '2', '4', '6', '8'})
elif p == 5:
    ans = sum(i for i, digit in enumerate(number, 1)
              if digit in {'0', '5'})

Note that I used the second (optional) argument of enumerate to start counting at one and I used a set instead of a tuple, because it has faster membership test (not a lot faster for small collections as this, but measurably faster). I also put it all into a generator expression and used sum.

---

Python has an official style-guide, PEP8, which recommends using lower_case for variables, so your counter should be called c, or better yet, something readable like digit_count. You can also factor out the 0 case to avoid having to check for it up to nine times:

else:
    slices = (int(number[i:]) % p for i in range(len(number)))
    digit_count = Counter(slices)

    # get the count of zeroes, with a default value of zero
    zeroes = digit_count.pop(0, 0)  
    ans = (zeroes * (zeroes + 1)) // 2    
    ans += sum((count * (count - 1)) // 2
               for count in digit_counts.values())

Answer 3

The issue that you're running into is two-fold. One, you're using bignum math instead of the much faster integer math. Since the original problem specifies that 2 <= p <= 10**9, integer math is certainly enough to deal with this issue (as Toby Speight's answer points out). The second is that you're enumerating all the (up to one million) subnumbers at the end of number.

I sign on to all of the comments in Graipher's answer, and add the following. Where he writes

slices = (int(number[i:]) % p for i in range(len(number)))

you should instead write

def generate_slices(number):
    mod = 0
    for digit in reversed(number):
        mod = (mod + int(digit)) % p
        yield mod
slices = generate_slices(number)

Answer 4

According to your guidances, I improved my code and it passed all tests.

from collections import Counter

number = input()
p = int(input())

ans = 0

if p == 2:
    ans = sum(i for i, digit in enumerate(number, 1)
              if digit in {'0', '2', '4', '6', '8'})
elif p == 5:
    ans = sum(i for i, digit in enumerate(number, 1)
              if digit in {'0', '5'})
else:
    def generate_slices(number):
        mod = 0
        power = 1
        for i in range(len(number) - 1, -1, -1):
            mod = (mod + int(number[i]) * power) % p
            power = (power * 10) % p
            yield mod
    slices = generate_slices(number)
    digit_count = Counter(slices)

    zeroes = digit_count.pop(0, 0)
    ans = (zeroes * (zeroes + 1)) // 2
    ans += sum((count * (count - 1)) // 2
               for count in digit_count.values())

print(ans)

Thank you all.

Answer 5

One thing we're missing is that we only care whether any substring is a multiple of the prime. That means we don't need to carry any state bigger than that prime.

Consider this algorithm:

1. initialize accumulator a with the first digit
2. reduce a modulo p (i.e, a = a % p)
3. if the result is zero, we found a multiple
4. append the next digit (a = a * 10 + next)
5. repeat 2-4 until we run out of digits

That still scales as O(n²) (as we need to run the loop once for each starting position), but we're now working with smaller intermediate numbers.

In practice, we'd probably combine steps 4 and 2 - a = (a * 10 + next) % p.

We can scale a little better by observing that any multiple (m1) directly followed by any other multiple (m2) also gives us a third one (m1m2), since multiplying 0 by a power of ten still gives 0. In general, a succession of n multiples will yield ½·n·(n-1) for our count (be careful not to double-count them, though).