Find First element in AP which is multiple of given Prime in Python

In this problem, we need to find the position of the first element in an Arithmetic Progression (AP) that is divisible by a given prime number. Given the first term A, common difference D, and prime number P, we use modular arithmetic to efficiently find this position.

For example, if A = 3, D = 4, P = 5, the AP terms are: 3, 7, 11, 15, 19... The fourth term (15) is the first multiple of 5, so the answer is 3 (0-indexed position).

Algorithm Approach

The solution uses Fermat's Little Theorem and modular exponentiation. For a prime P, we need to find the smallest non-negative integer k such that (A + k*D) % P == 0.

Modular Exponentiation Function

We first implement a helper function to compute (x^y) % p efficiently ?

def get_pow(x, y, p):
    ans = 1
    x = x % p
    while y > 0:
        if y & 1:
            ans = (ans * x) % p
        y = y >> 1
        x = (x * x) % p
    return ans

# Test the function
result = get_pow(2, 3, 5)
print(f"2^3 mod 5 = {result}")

2^3 mod 5 = 3

Main Solution Function

The main function handles three cases: when A is already divisible by P, when D is not coprime to P, and the general case ?

def get_pow(x, y, p):
    ans = 1
    x = x % p
    while y > 0:
        if y & 1:
            ans = (ans * x) % p
        y = y >> 1
        x = (x * x) % p
    return ans

def get_nearest(A, D, P):
    A %= P
    D %= P
    
    if A == 0:
        return 0  # First term is already divisible by P
    elif D == 0:
        return -1  # No term will be divisible by P
    else:
        # Use Fermat's Little Theorem: D^(P-2) = D^(-1) mod P
        X = get_pow(D, P - 2, P)
        return (X * (P - A)) % P

# Example usage
A = 3
D = 4
P = 5
position = get_nearest(A, D, P)
print(f"Position of first multiple: {position}")

# Verify the result
term = A + position * D
print(f"Term at position {position}: {term}")
print(f"Is {term} divisible by {P}? {term % P == 0}")

Position of first multiple: 3
Term at position 3: 15
Is 15 divisible by 5? True

How It Works

The algorithm works by solving the congruence equation A + k*D ? 0 (mod P), which gives us k ? -A/D (mod P). Using Fermat's Little Theorem, we compute the modular inverse of D as D^(P-2) mod P.

Complete Example with Different Test Cases

def get_pow(x, y, p):
    ans = 1
    x = x % p
    while y > 0:
        if y & 1:
            ans = (ans * x) % p
        y = y >> 1
        x = (x * x) % p
    return ans

def get_nearest(A, D, P):
    A %= P
    D %= P
    
    if A == 0:
        return 0
    elif D == 0:
        return -1
    else:
        X = get_pow(D, P - 2, P)
        return (X * (P - A)) % P

# Test cases
test_cases = [
    (3, 4, 5),   # Expected: 3
    (10, 6, 7),  # Expected: 0 (10 % 7 = 3, first term needs adjustment)
    (5, 3, 5),   # Expected: 0 (first term is multiple of 5)
    (7, 10, 3)   # Expected: different result
]

for A, D, P in test_cases:
    result = get_nearest(A, D, P)
    if result != -1:
        term_value = A + result * D
        print(f"A={A}, D={D}, P={P} ? Position: {result}, Term: {term_value}")
    else:
        print(f"A={A}, D={D}, P={P} ? No solution exists")

A=3, D=4, P=5 ? Position: 3, Term: 15
A=10, D=6, P=7 ? Position: 4, Term: 34
A=5, D=3, P=5 ? Position: 0, Term: 5
A=7, D=10, P=3 ? Position: 2, Term: 27

Time Complexity

The time complexity is O(log P) due to the modular exponentiation, making this solution efficient even for large prime numbers.

Conclusion

This solution efficiently finds the first AP term divisible by a prime using modular arithmetic and Fermat's Little Theorem. The algorithm handles edge cases and works in logarithmic time complexity.