[백준] 1644. 소수의 연속합 - 파이썬

[Gold III]

 

https://www.acmicpc.net/problem/1644

1644번: 소수의 연속합

 

풀이 1

보자마자, Prefix Sum 문제라고 생각했다.

먼저 에라토스테네스의 체를 이용해 입력값인 N까지의 수 중 소수를 구하고,

그 소수들의 리스트를 이용해 prefix sum 리스트를 만든 후,

Two Pointer 기법을 이용해 부분합의 값이 N과 동일한 지 확인하는 방법.

시간은 충분했으나,

MLE. (Memory Limit Exceed.)

N = int(input())

nums = {i: 0 for i in range(2, N+1)}

# 에라토스테네스의 체
for i in nums:
    if nums[i] > 0:
        continue
    elif i > N:
        break

    k = 1
    while True:
        k += 1
        j = i*k
        if j > N:
            break
        nums[j] += 1

primes = []
for i in nums:
    if nums[i] == 0:
        primes.append(i)

# prefix sum of primes
sum_prime = [0, primes[0]]
for i in range(1, len(primes)):
    sum_prime.append(sum_prime[i] + primes[i])

# two pointer of index
p1 = 0
p2 = 0
result = 0
while True:
    S = sum_prime[p2] - sum_prime[p1]

    if S == N:
        p2 += 1
        result += 1
    elif S < N:
        p2 += 1
    elif S > N:
        p1 += 1

    if p1 >= len(sum_prime) or p2 >= len(sum_prime):
        break

print(result)

 

풀이 2

메모리 사용을 줄이기 위해, sum_prime 리스트를 따로 만들지 않고, primes 리스트에서 List Slicing으로 부분합을 처리하고,

처음에 에라토스테네스의 체로 소수를 판정할 때, Dictionary가 아닌, True&Flase를 저장하는 배열로 처리하였다.

AC.

N = int(input())

nums = [True] * (N+1)
nums[0] = False
nums[1] = False

# 에라토스테네스의 체
for i in range(2, N+1):
    if not nums[i]:
        continue
    elif i > N:
        break

    k = 1
    while True:
        k += 1
        j = i*k
        if j > N:
            break
        nums[j] = False

primes = []
for i, flag in enumerate(nums):
    if flag:
        primes.append(i)

# two pointer of index
p1 = 0
p2 = 0
result = 0
while True:
    S = sum(primes[p1:p2])

    if S == N:
        p2 += 1
        result += 1
    elif S < N:
        p2 += 1
    elif S > N:
        p1 += 1

    if p1 > len(primes) or p2 > len(primes):
        break

print(result)