[BOJ][Python] 백준 14317번 - Sherlock and Watson Gym Secrets (Large)

 

문제 링크: https://www.acmicpc.net/problem/14317

---

문제 풀이

정수론, 포함 배제의 원리, 분할 정복을 이용한 거듭제곱

 

$N \leq 10^{18}$인 $i$, $j$가 있을 때 $(i^{A} + j^{B}) \pmod K = 0$와 $i \neq j$를 만족하는 순서 쌍 ($i$, $j$)의 개수를 구하는 문제이다. 굳이 위에 $10^{18}$을 언급하는 이유는 나이브한 $O(N^{2})$ 풀이는 쳐다도 볼 수가 없다는 뜻이다.

 

여기서 주목해야 하는건 $K$의 범위다. $10^{5}$보다 작거나 같으므로 이를 이용해보자. 먼저, '$N$보다 작거나 같은 $i$가 $K$로 나누었을 때 나머지가 $r$인 값은 총 몇 개 있는가?' 라는 해답부터 생각하면 된다. 당연하겠지만 $0 \leq r < K$인 모든 값을 뜻한다. $N$을 $K$로 나누었을 때의 몫을 $p$, $N$을 $K$로 나누었을 때의 나머지를 $q$라고 하면, $r = 0$ 또는 $q < r < K$인 것들은 각각 $p$개 있고, $1 \leq r \leq q$인 것들은 각각 $p+1$개 있다. 개수들을 배열에다가 채워놓자. 나는 배열의 인덱스를 $r$로 정하고 개수를 채워나갔다. 

 

그러면 $r^{A} \bmod K = u \ (0 \leq u < K) $도 생각해볼 수 있다. 여기서 구하고자 하는건 $i \bmod K \ (i \leq N)$를 했을 때 나머지가 $r$인 값들을 $A$제곱하고 $K$로 나눈 나머지가 $u$인 값의 개수를 구하자는 의미다. $i \bmod K$를 했을 때 나머지가 $r$인 값은 우리가 위에서 발견했으니, $r^{A} \bmod K$한 값 $x$를 구하고 $\text{acnt}$라는 $K$ 크기의 배열을 생성한 후 $\text{acnt[x]}$에 아까 구한 배열을 이용해서 더해주면 된다. $r^{B} \bmod K = u \ (0 \leq u < K) $도 마찬가지로 구해준다.

 

그러면 원래 식 $(i^{A} + j^{B}) \pmod K = 0$으로 돌아와서 $i^{A} \bmod K = u$라면, $j^{B} \bmod K = (K-u) \bmod K \ (0 \leq u < K) $면 성립한다. 이건 위에서 다 구해줬으니 $u$에 관해서 반복문을 돌려주고, 마지막으로 $i = j$한 경우만 빼주자. 즉 $(i^{A} + i^{B}) \pmod K = 0$인 값들만 빼주면 되므로 위에서 구한 배열로 활용하자.

 

경우의 수는 $10^{9}+7$로 나눈 나머지 값을 구하라고 했으니 나머지 연산을 잘 활용하여 답을 구하자.

코드