문제 개요
등차수열(Arithmetic Progression, AP)의 첫째 항 A와 공차 D가 주어져 있고, 소수 P가 하나 더 주어졌다고 가정해 보겠습니다. 이때 구해야 하는 것은 해당 등차수열에서 처음으로 소수 P의 배수가 되는 항의 위치입니다.
예를 들어 입력이 A = 3, D = 4, P = 5라면 결과는 3이 됩니다. 네 번째 항이 소수 5의 배수이기 때문인데, 실제로 항들을 확인해 보면 다음과 같습니다.
- 첫째 항 = 3
- 둘째 항 = 3 + 4 = 7
- 셋째 항 = 3 + 2×4 = 11
- 넷째 항 = 3 + 3×4 = 15 → 5의 배수
풀이 접근 방식
k번째 항은 A + (k−1)×D 형태로 표현됩니다. 따라서 A + (k−1)×D ≡ 0 (mod P)을 만족하는 가장 작은 k를 찾으면 되며, 이 식을 정리하면 다음과 같습니다.
k − 1 ≡ −A × D⁻¹ (mod P)
여기서 D⁻¹은 P에 대한 D의 모듈러 역원(modular inverse)이며, 페르마의 소정리(Fermat's Little Theorem)에 의해 D^(P−2) mod P로 계산할 수 있습니다. 모듈러 거듭제곱은 반복 제곱 기법을 사용하면 O(log P) 시간 안에 효율적으로 구할 수 있습니다.
알고리즘 단계
- 모듈러 거듭제곱을 계산하는 함수 get_pow(x, y, p)를 정의합니다.
- ans를 1로 초기화하고, x를 x mod p 값으로 줄입니다.
- y > 0인 동안 반복합니다. y의 최하위 비트가 1이면 ans에 x를 곱한 뒤 p로 나눈 나머지를 저장하고, y를 절반으로 줄이면서 x는 x² mod p로 갱신합니다.
- 메인 로직에서는 먼저 A := A mod P, D := D mod P로 값을 줄입니다.
- A가 0이라면 첫째 항이 이미 P의 배수이므로 0을 반환합니다.
- D가 0이라면 수열의 모든 항이 동일하므로 P의 배수가 존재하지 않아 -1을 반환합니다.
- 그 외의 경우 X = get_pow(D, P − 2, P)로 모듈러 역원을 구한 후, (X × (P − A)) mod P를 반환합니다.
구현 예제
아래 파이썬 코드를 통해 실제 동작 과정을 확인해 볼 수 있습니다.
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
A = 3
D = 4
P = 5
print(get_nearest(A, D, P))
입력
A = 3, D = 4, P = 5
출력
3
복잡도 분석
시간 복잡도는 모듈러 거듭제곱 연산이 지배하므로 O(log P)이며, 추가 메모리를 거의 사용하지 않으므로 공간 복잡도는 O(1)입니다. 덕분에 P가 매우 큰 소수일 때도 빠르게 답을 구할 수 있다는 장점이 있습니다.