ABC 228 E Integer Sequence Fair 풀이

오늘 대회 중에 아쉽게 못 풀었던 문제입니다. 

페르마의 소정리를 안다면 어렵지 않게 풀 수 있고, 좋은 문제인지는 잘 모르겠지만 개인적으로 이 문제를 풀며 몇 가지를 배우게 되어 정리할 겸 풀이를 작성하려 합니다.

 

문제

https://atcoder.jp/contests/abc228/tasks/abc228_e

E - Integer Sequence Fair

1 이상 K 이하의 자연수가 N개 있는 모든 배열에 1 이상 M 이하의 가치를 매기는 경우의 수를 $Q = 998244353$으로 나눈 나머지를 구하는 문제입니다. 링크된 문제의 예제를 통해 문제를 쉽게 이해할 수 있습니다.

 

풀이

조건에 맞는 배열이 몇 개인지부터 알아봅시다.

배열의 수는 중복을 허용하며 K개의 자연수 중 N개를 뽑는 경우의 수와 같으므로 조건에 맞는 배열은 총 $K^N$개입니다. 각 배열에 가치를 매기는 경우의 수는 역시 중복을 허용하여 M개의 자연수 중 배열의 수만큼 뽑는 경우의 수와 같으므로 전체 경우의 수는 $M^{K^N}$입니다. 따라서 구해야 하는 정답은 $M^{K^N}\pmod{Q}$입니다. 

 

$1\leq N, M, K\leq 10^{18}$이므로 당연히 무작정 구할 수는 없습니다. 페르마의 소정리에 의하면 정수 $a$와 서로소인 소수 $p$에 대해 $a^{p-1}\equiv 1 \pmod p$ 가 항상 성립합니다. Q는 소수이므로 M이 Q의 배수만 아니라면 $M^{K^N}$에 $M^{Q-1}$을 곱하건 나누건 Q로 나눈 나머지는 변하지 않습니다. 따라서 $r = K^N\pmod{Q-1}$이라고 할 때 $M^r\pmod{Q}$를 구하면 정답을 알 수 있을 것 같습니다.

 

대회 중에는 여기까지만 생각했고, 맞왜틀을 외치다 대회가 끝났습니다.

위 풀이는 K가 $Q-1$의 배수라면 1을 출력한다는 문제가 있습니다.

M이 Q의 배수라면 $M^{Q-1}\equiv 1 \pmod Q$가 성립하지 않으며, 무조건 0을 출력해야 합니다. 

이러한 처리만 추가하면 문제를 풀 수 있습니다.

 

코드

https://atcoder.jp/contests/abc228/submissions/27403485

Submission #27403485 - TOYOTA SYSTEMS Programming Contest 2021(AtCoder Beginner Contest 228)

다른 분의 정답 코드를 대신 첨부합니다.

파이썬의 pow 함수는 거듭제곱을 계산하는 함수이며, pow(a, b, M)과 같은 형태로 호출하면 $a^b$를 M으로 나눈 값을 반환합니다. 시간복잡도는 $O(\log b)$입니다.

 

+

정수 a가 소수 P의 배수라면 법 P에서 a의 모듈러 역원은 존재하지 않습니다.

모듈러 역원은 유일합니다. 즉 확장 유클리드 알고리즘을 써도, 페르마의 소정리를 써도 동일한 역원을 얻습니다.

 

잘못된 내용이 있다면 가차없이 지적해주세요