[백준 1644번] 소수의 연속합 (Python3)

기본적인 알고리즘은 일단 에라토스테네스의 체 방식으로 소수 리스트를 구한 뒤 소수의 연속합을 구하면 된다.

 

처음에는 이중 for문을 이용해서 돌렸다. 

 

import sys
input = sys.stdin.readline
from collections import deque

N = int(input())

DP = [True]*(N+1)

DP[0] = DP[1] = False

prime = deque()
for i in range(1,N+1):
  if DP[i]:
    prime.append(i)
    c = 2
    while True:
      if i*c > N:
        break
      DP[i*c] = False
      c+=1

cnt = 0
l = len(prime)

for i in range(l):
  SUM = 0
  No = i
  while True:
    if SUM == N:
      cnt+=1
      break
    if SUM > N or No >= l:
      break
    SUM += prime[No]
    No += 1
print(cnt)

 

그러자 pypy에서는 잘 통과가 되었는데 python에서는 시간초과로 통과되지 못하였다.

 

그러다가 지난번에 부분합 문제를 풀면서 배운 "투 포인터 알고리즘"이 생각났다.

 

아래는 투 포인터 알고리즘을 적용한 코드이다.

 

import sys
input = sys.stdin.readline
from collections import deque

N = int(input())

DP = [True]*(N+1)

DP[0] = DP[1] = False

prime = deque()
for i in range(1,N+1):
  if DP[i]:
    prime.append(i)
    c = 2
    while True:
      if i*c > N:
        break
      DP[i*c] = False
      c+=1

cnt = SUM = start = end = 0
while True:
  if SUM < N:
    if end >= len(prime):
      break
    SUM += prime[end]
    end += 1
  elif SUM==N:
    cnt += 1
    if end >= len(prime):
      break
    SUM += prime[end]
    end += 1
  elif SUM > N:
    SUM -= prime[start]
    start += 1

print(cnt)

 

이 방식을 이용하니 파이썬에서도 통과가 되었다.

그런데 문제는 pypy에서는 오히려 이중 for문보다 시간이 더 오래 걸렸다.

왜 이런 결과가 나온건지는 나도 모르겠다...

 

그리고 에라토스테네스의 체 방식을 이용할 때 다른 사람들 코드를 참고하니 1부터 N까지 다 돌리는게 아니라 N의 제곱근 까지만 돌려봐도 된다는 사실을 깨달았다. 다음에는 이 방식을 이용해서 시간을 좀 더 줄여야겠다.

 

 

오늘의 교훈) 수학이 중요하다