[Baekjoon] 7512/연속하는 소수의 합/Python/파이썬/누적 합/슬라이딩 윈도우

井·2025년 2월 4일

PrefixSum python slidingwindow 누적합 백준 소수판정 수학 슬라이딩윈도우 알고리즘 에라토스테네스의체 정수론 파이썬

문제풀이

목록 보기

29/56

💡문제

숫자 m개가 주어지며, 각 숫자를 ni라고 한다. (1 ≤ i ≤ m)

이때, 모든 i에 대해서, 연속하는 소수 ni개의 합으로 나타낼 수 있는 가장 작은 소수를 찾는 프로그램을 작성하시오.

예를 들어, m = 2, n1 = 3, n2 = 5인 경우에 정답은 83이다. (83 = 23 + 29 + 31 = 11 + 13 + 17 + 19 + 23)

입력

첫째 줄에 테스트 케이스의 개수가 주어진다. 각 테스트 케이스의 첫째 줄에는 1 ≤ m ≤ 10 이 주어진다. 둘째 줄에는 ni가 주어진다. (1 ≤ ni ≤ 104)

항상 정답은 107보다 작다.

출력

각 테스트 케이스마다 "Scenario i:"를 출력하고, 둘째 줄에 정답을 출력한다.

예제입력

예제출력

Scenario 1:
83

Scenario 2:
311

📖내가 실패한 Code

import sys

''''
숫자 m개에 대해 모든 i에 대해 
'''

MAX_NUMBER = 10**7


def make_sieve_of_eratosthenes():
    is_prim = [1] * (MAX_NUMBER + 1)
    is_prim[0] = is_prim[1] = 0

    for num in range(2, int(MAX_NUMBER**0.5) + 1):
        if is_prim[num]:
            for i in range(num * num, MAX_NUMBER + 1, num):
                is_prim[i] = 0

    prime_list = [i for i, flag in enumerate(is_prim) if flag]
    return is_prim, prime_list


class ConsecutivePrimeSum:
    is_prim, prime_list = make_sieve_of_eratosthenes()

    def __init__(self, number):

        start = 0
        end = number - 1
        prime_sum = sum(ConsecutivePrimeSum.prime_list[:end+1])

        while not ConsecutivePrimeSum.is_prim[prime_sum]:
            prime_sum -= ConsecutivePrimeSum.prime_list[start]
            start += 1
            end += 1
            prime_sum += ConsecutivePrimeSum.prime_list[end]

        self.start = start
        self.end = end
        self.prime_sum = prime_sum

    def next_prime_sum(self):
        while True:
            self.prime_sum -= ConsecutivePrimeSum.prime_list[self.start]
            self.start += 1
            self.end += 1
            self.prime_sum += ConsecutivePrimeSum.prime_list[self.end]

            if ConsecutivePrimeSum.is_prim[self.prime_sum]:
                break


def find_same_sum(objs):
    while True:
        objs.sort(key=lambda x: x.prime_sum)
        objs[0].next_prime_sum()
        if all(obj.prime_sum == objs[0].prime_sum for obj in objs):
            break

    return objs[0].prime_sum


def main():
    inputs = map(str.split, sys.stdin.read().splitlines())

    for test_case in range(1, int(next(inputs)[0])+1):
        N = int(next(inputs)[0])
        array = next(inputs)
        result = find_same_sum([ConsecutivePrimeSum(int(array[num])) for num in range(N)])
        sys.stdout.write(f"Scenario {test_case}:\n{result}\n\n")


if __name__ == '__main__':
    main()

📖내가 작성한 Code

import sys

''''
와 어려워
'''

MAX_NUMBER = 10**7


def make_sieve_of_eratosthenes():
    is_prim = [1] * (MAX_NUMBER + 1)
    is_prim[0] = is_prim[1] = 0

    for num in range(2, int(MAX_NUMBER**0.5) + 1):
        if is_prim[num]:
            for i in range(num * num, MAX_NUMBER + 1, num):
                is_prim[i] = 0

    prime_list = [i for i, flag in enumerate(is_prim) if flag]
    return is_prim, prime_list


class ConsecutivePrimeSum:
    is_prim, prime_list = make_sieve_of_eratosthenes()

    def __init__(self, number):

        start = 0
        end = number - 1
        prime_sum = sum(ConsecutivePrimeSum.prime_list[:end+1])

        while not ConsecutivePrimeSum.is_prim[prime_sum]:
            prime_sum -= ConsecutivePrimeSum.prime_list[start]
            start += 1
            end += 1
            prime_sum += ConsecutivePrimeSum.prime_list[end]

        self.start = start
        self.end = end
        self.prime_sum = prime_sum

    def next_prime_sum(self):
        while True:
            self.prime_sum -= ConsecutivePrimeSum.prime_list[self.start]
            self.start += 1
            self.end += 1
            self.prime_sum += ConsecutivePrimeSum.prime_list[self.end]

            if ConsecutivePrimeSum.is_prim[self.prime_sum]:
                break


def find_same_sum(objs):
    while not all(obj.prime_sum == objs[0].prime_sum for obj in objs):
        objs.sort(key=lambda x: x.prime_sum)
        objs[0].next_prime_sum()

    return objs[0].prime_sum


def main():
    inputs = map(str.split, sys.stdin.read().splitlines())

    for test_case in range(1, int(next(inputs)[0])+1):
        N = int(next(inputs)[0])
        array = next(inputs)
        result = find_same_sum([ConsecutivePrimeSum(int(array[num])) for num in range(N)])
        sys.stdout.write(f"Scenario {test_case}:\n{result}\n\n")


if __name__ == '__main__':
    main()

✍️풀이과정

슬라이딩 윈도우에 익숙해지기 위해서 잡은 문제인데 이렇게 어려울 줄은 몰랐음.

감이 안잡혔음. 단계별로 생각해봄.

10^7보고 에라토스테네스의 체를 사용해야겠다는 생각을 하여서 만들어봄.

그런데 슬라이딩윈도우를 이용해야하기 때문에, 그 소수들을 리스트에 따로 넣어줌

각각 시작점과 끝점, 그리고 덧셈 값이 필요해서, 리스트로 만들까 하다가 내가 코드를 못알아 볼꺼 같아서 객체로 만들어 이어나감.

그래도 너무 길어져서, 정말 간단한 실수들을 하고 말았음. 슬라이싱을 실수하는건 진짜 오랫만...

이후 다음 누적합 찾는 곳에서 논리적 실수를 저지름. 먼저 누적합들이 같은지 체크하고 다음 누적합을 찾았어야했었음.

이런 방법으로 정답을 찾아가는데 3시간은 걸린 거 같음. 오랫만에 알고리즘을 깊게 풀고 머리쓰는 시간을 가진듯.

열심히 해서인지, 전국에서 속도로 2등할 줄은 몰랐음. 저번에도 그렇고 아무 기대없이 봤다가 이러면 기분이 은근히 좋음.

🧠 코드 리뷰

코드 주석 및 변수 명: 좀 더 자세한 주석과 직관적인 변수명을 사용하면 유지보수나 타인이 코드를 이해하는 데 도움이 됩니다.
에러 핸들링: 슬라이딩 윈도우 이동 시 인덱스 범위가 초과할 가능성에 대한 검증(예외 처리 등)을 고려하면 안정성이 높아집니다.
입력 처리 방식: 입력 데이터를 보다 명시적으로 처리하는 방식(예: 리스트로 변환 후 인덱스 접근)도 고려해보세요.

🛠️AI 개선 코드

import sys

MAX_NUMBER = 10**7

def make_sieve_of_eratosthenes():
    """
    에라토스테네스의 체를 이용해 0부터 MAX_NUMBER까지의 소수를 판별한다.
    반환:
        is_prime: 인덱스가 소수이면 True, 아니면 False 인 리스트
        prime_list: 0부터 MAX_NUMBER 이하의 모든 소수를 담은 리스트
    """
    is_prime = [True] * (MAX_NUMBER + 1)
    is_prime[0] = False
    is_prime[1] = False

    # 2부터 sqrt(MAX_NUMBER)까지 배수 지우기
    limit = int(MAX_NUMBER**0.5)
    for num in range(2, limit + 1):
        if is_prime[num]:
            for multiple in range(num * num, MAX_NUMBER + 1, num):
                is_prime[multiple] = False

    # 소수 리스트 생성
    prime_list = [i for i, flag in enumerate(is_prime) if flag]
    return is_prime, prime_list

class ConsecutivePrimeSum:
    """
    특정 개수(n)만큼 연속되는 소수의 합을 관리하는 클래스.
    """
    is_prime, prime_list = make_sieve_of_eratosthenes()

    def __init__(self, n):
        """
        n개의 연속되는 소수 합이 '소수'가 되는 가장 작은 합을 찾는다.
        start와 end는 prime_list 내 인덱스, prime_sum은 소수 합 값을 저장.
        """
        self.start = 0
        self.end = n - 1
        # 처음 n개 소수의 합
        self.prime_sum = sum(ConsecutivePrimeSum.prime_list[self.start : self.end + 1])

        # 초기 합이 소수가 아닐 경우, 소수가 될 때까지 한 칸씩 윈도우를 이동
        while not ConsecutivePrimeSum.is_prime[self.prime_sum]:
            self.prime_sum -= ConsecutivePrimeSum.prime_list[self.start]
            self.start += 1
            self.end += 1
            # 인덱스 범위를 넘어갈 위험이 있으나, 문제에서 정답 < 10^7 보장
            self.prime_sum += ConsecutivePrimeSum.prime_list[self.end]

    def next_prime_sum(self):
        """
        연속 합의 '다음' 소수 후보로 이동한다.
        현재 윈도우를 한 칸 오른쪽으로 밀고(슬라이딩),
        그 합이 소수가 될 때까지 반복.
        """
        while True:
            self.prime_sum -= ConsecutivePrimeSum.prime_list[self.start]
            self.start += 1
            self.end += 1
            self.prime_sum += ConsecutivePrimeSum.prime_list[self.end]

            if ConsecutivePrimeSum.is_prime[self.prime_sum]:
                break

def find_same_sum(objs):
    """
    여러 ConsecutivePrimeSum 객체가 모두 같은 prime_sum을 가질 때까지,
    매번 prime_sum이 가장 작은 객체만 다음 candidate로 이동시킨다.
    모든 prime_sum이 동일해지면 그 값을 반환한다.
    """
    while True:
        # 먼저 모두 동일한지 확인
        if all(obj.prime_sum == objs[0].prime_sum for obj in objs):
            return objs[0].prime_sum

        # 같지 않다면, 현재 가장 작은 값을 가진 객체를 찾는다
        objs.sort(key=lambda x: x.prime_sum)
        # 그 객체를 다음 소수 후보로 이동
        objs[0].next_prime_sum()

def main():
    input_lines = sys.stdin.read().splitlines()
    # 첫 줄: 테스트 케이스 개수
    t = int(input_lines[0])
    idx = 1  # 현재 읽을 줄의 인덱스

    for scenario_num in range(1, t + 1):
        # 이 테스트 케이스에서 m(연속 소수 개수의 개수)
        m = int(input_lines[idx])
        idx += 1
        # m개의 숫자 (연속 소수 개수들)
        nums = list(map(int, input_lines[idx].split()))
        idx += 1

        # 각 연속 소수 개수에 대해 객체를 생성
        objects = [ConsecutivePrimeSum(n) for n in nums]
        result = find_same_sum(objects)

        # 결과 출력
        sys.stdout.write(f"Scenario {scenario_num}:\n{result}\n\n")

if __name__ == "__main__":
    main()