BOJ 5615 - 아파트 임대 (Python3)
밀러-라빈과 폴라드-로 알고리즘을 같이 쓰는 경우가 많아서 간과하는 경우는 많지만, 밀러-라빈 하나만으로도 이미 충분히 쓸모가 있는 알고리즘이다.
그도 그럴 것이, 어떤 수가 소수인지 아닌지 로그 시간만에 판별 가능한 알고리즘이 그렇게 적지는 않다!
밀러-라빈 단독으로 사용하는 대표적인 문제 중 하나가 바로 이 문제이다.
그럼 본격적으로 풀이에 돌입해 보자.
문제
동규부동산에서 아파트를 임대하고 있다. 아파트의 방은 아래 그림과 같이 면적이 2xy+x+y이다. (x와 y는 양의 정수)
동규부동산의 카탈로그에는 아파트의 면적이 오름차순으로 적혀져 있지만, 이 중 일부는 있을 수 없는 크기의 아파트이다. 만약, 이런 크기의 아파트를 임대하겠다고 말하면, 동규는 꽝! 이라고 외치면서, 수수료만 떼어간다.
동규부동산의 카탈로그에 적힌 아파트의 면적이 주어졌을 때, 있을 수 없는 크기의 아파트의 수를 구하는 프로그램을 작성하시오.
입력
첫째 줄에 아파트의 면적의 수 N이 주어진다. 다음 줄부터 N개 줄에 카탈로그에 적혀있는 순서대로 면적이 주어진다. N은 100,000 이하이고 면적은 \(2^{31}-1\)이하인 양의 정수이다.
예제 입력)
10
4
7
9
10
12
13
16
17
19
20
출력
첫째 줄에 있을 수 없는 아파트 면적의 수를 출력한다.
예제 출력)
2
내 코드
from random import randrange
from sys import setrecursionlimit, stdin
pli = [2,13,23,1662803]
input = stdin.readline
# Miller-Rabin
def power(x,y,p):
if y<2: return (x**y)%p
else:
d = y//2
if y%2 == 0: return (power(x,d,p)**2)%p
else: return (x*(power(x,d,p))**2)%p
def miller_rabin(n,a):
r,d = 0,n-1
while d%2 == 0: r += 1; d //= 2
x = power(a,d,n)
if x == 1 or x+1 == n: return True
for i in range(0, r-1):
x = power(x,2,n)
if x+1 == n: return True
return False
def isprime(n):
if n in pli: return True
if n == 1 or n%2 == 0: return False
for p in pli:
if not miller_rabin(n,p): return False
return True
# Answer Code
cnt = 0
for _ in range(int(input())):
if isprime(1+2*int(input())): cnt += 1
print(cnt)
마지막 네 줄을 제외하고는 모두 밀러-라빈을 구현하는 부분이다.
먼저 집의 갯수가 10만개이기 때문에, 입력을 빠르게 하는 명령어인 sys.stdin.readline과 좀 더 빠른 언어인 PyPy3을 사용하지 않는다면 얼마나 코드를 아름답게 짰더라도 시간 초과가 나올 것이다.
문제 풀이에 메인이 되는 것은 소인수분해이다. 집의 넓이에 두 배를 한 값은 4xy+2x+2y로, 이는 (2x+1)(2y+1)-1로 인수분해가 된다.
그 말인즉슨, 집의 넓이를 S라고 두었을 때 2S+1=(2x+1)(2y+1)을 만족한다는 것이다. 그렇지 않다면? 집이 아니다!
따라서 2S+1이 소수인지 아닌지만 판정하면 된다. 그렇다. 밀러-라빈 알고리즘을 사용하는 것이다.
이로서 5651번의 풀이를 마친다.