[PYTHON] Calculation problem using mod

at first

When the following problems were asked in AtCoder, I couldn't help at all, so I will summarize the theory and its implementation.

Multiple of 2019 ( AtCoder Beginner Contest 164; Problem D )

Problem statement You will be given the string S, which consists only of numbers from 1 to 9. A set of integers that satisfy the following conditions(i,j) (1≤i≤j≤|S|)Find the total number of.

Condition: If you look at the i-th to j-th letters of S as decimal integers, this integer is a multiple of 2019.

Constraint 1≤|S|≤200000 S is a string consisting only of numbers from 1 to 9

theory

It is possible to calculate by stacking for loops in a straightforward manner, but in the case of AtCoder, it is necessary to devise because it exceeds the upper limit of the calculation time.

This kind of problem of calculating multiples (or the problem of thinking about radix?) Can be used to write efficient programs by algorithms that make good use of mod calculations.

mod basic formula

a \equiv b \pmod{ 2019 } \\
c \equiv d \pmod{ 2019 }

When

a + c \equiv b + d \pmod{ 2019 } \\
a - c \equiv b - d \pmod{ 2019 } \\
a * c \equiv b * d \pmod{ 2019 }

Is established. In other words, the sum, difference, and product formulas in ordinary four arithmetic operations can be applied to mod calculations.

By the way, the following rules also hold for quotients.

If $ ab \ equiv ac $ and $ a $ and $ p $ are relatively prime, then $ b \ equiv c $

Application to this problem

The given string $ S $ is

S = s_n s_{n-1} \dots s_2 s_1

Suppose that In other words, contrary to the problem, consider setting $ i <j $. What is the set $ (i, j) $ that meets the conditions in the problem?

N_{ij} = 10^{i-j} s_i + 10^{i-j-1} s_{i+1} + \dots + 10 s_{j+1} + s_j

Is a multiple of 2019. It is required to count the set $ (i, j) $ of this condition. First, define the following $ r_i $ and $ t_i $.

t_i = 10^{i-1} \\
r_i = t_i s_i + t_{i-1} s_{i-1} + \dots + t_1 s_1

This $ r_i $ can also be written as follows.

r_i = t_i s_i + r_{i-1}

At this time, the following equation holds.


r_{i} \equiv r_{j-1} \Rightarrow N_{ij} \equiv 0 \pmod{ 2019 } 

This proof is shown below. The following relationship holds between $ r_i $ and $ N_ {ij} $.

r_{i} - r_{j-1} = N_{ij} * 10^{j-1} 

If $ r_ {i} \ equiv r_ {j-1} $

r_{i} - r_{j-1} = N_{ij} * 10^{j-1} \equiv 0 \pmod{ 2019 } 

Because it becomes

N_{ij} \equiv 0 \pmod{ 2019 } 

Or

10^{j-1} \equiv 0 \pmod{ 2019 } 

Should hold. But since $ 10 ^ {j-1} = 2 ^ {j-1} * 5 ^ {j-1} $, $ 2019 $ and $ 10 ^ {j-1} $ are clearly disjoint, the former $ N_ {ij} \ equiv 0 $ will hold.

algorithm

r_i \equiv R_i \pmod{ 2019 } \\
s_i \equiv S_i \pmod{ 2019 } \\
t_i \equiv T_i \pmod{ 2019 }

However, $ R_i, S_i, T_i $ represent the remainder, which is $ 0 \ le R_i, S_i, T_i \ le 2018 $. Using these characters, the remainder of each $ r_i $ can be calculated continuously as follows.

\begin{align}
 r_i &= t_i s_i + r_{i-1} \\
& \equiv T_i S_i + R_{i-1} = R_i \pmod{ 2019 } \\
\end{align}

You can use this method to calculate the remainder for each $ r_i $ for 2019, and count the combinations of $ r_i $ with the same remainder to calculate the answer.

Implementation

By using the above algorithm, it is possible to calculate without overlapping for loops.

s = input()

d = [ int( 0 ) ] * 2019
d[ 0 ] = int( 1 )
r = int( 0 )
t = int( 1 )

for si in reversed( s ):
#    print( si )

    r = int( si ) * t + r
    r = r % 2019

    t = t * 10
    t = t % 2019

    d[ r ] = d[ r ] + 1

ans = int( 0 )
for di in d:
    ans = ans + di * ( di - 1 ) // 2

print( ans )

Recommended Posts

Calculation problem using mod
Calculation time measurement using maf
AtCoder ABC156 Problem D Bouquet mod calculation and permutations, etc.
[Mathematical optimization problem] Linear programming using PuLP
[Scientific / technical calculation by Python] Solving (generalized) eigenvalue problem using numpy / scipy, using library