백준 9020 FP 적용해서 풀기

2022 Dec 07 in Algorithm on Algorithm, Python, Fp

백준 9020 골드바흐의 추측 문제를 FP 적용해서 풀어보았다.

• 백준 9020 FP 적용해서 풀기
  • 기존 풀이
  • 함수형 패러다임
    • 에라토네스 체 구현
    • 본문
  • 풀이

백준 9020 문제에 함수형 패러다임을 적용해서 풀어보았다.

기존 풀이

누군가 내 기존 풀이를 봤다는 알림을 받아서 내 코드를 확인해보니 마음에 들지 않았다.

def solution(open = open):
    che = [False, False, True, True] + [False, True] * 4998  # 에라토스테네스의 체
    # 0과 1은 소수가 아니므로 False, 2와 3부터 소수이므로 True,
    # 이후 짝수만 False, 홀수만 True로 초기화한다.
    for i, is_prime in enumerate(che[3:int(10000 ** .5) + 1:2], start=1):
        # 3부터 10000(문제에서의 최댓값)의 제곱근까지 홀수만 검사한다.
        if not is_prime:
            # 홀수가 아니면 넘긴다.
            continue
        prime = i * 2 + 1
        # 인덱스 값에서 원래 수를 구한다.
        che[prime * prime::2 * prime] = [False] * ((10000 - prime * prime) // (2 * prime) + 1)
        # prime의 제곱부터 2 * prime 간격으로 False로 초기화한다.
        # 제곱부터인 이유는 그보다 작은 prime의 배수는 이미 prime보다
        # 작은 수와의 공배수이므로 False로 초기화되어 있기 때문이다.
        # 2 * prime 간격으로 초기화하는 이유는 짝수는 이미 False로 초기화되어 있기 때문이다.
    input = iter(open(0).read().split("\n")).__next__
    # open(0)은 표준 입력을 의미한다. 알면 왜 쓰는지 알테고 모르면 그냥 그렇구나 하면 된다.
    # open(0).read().split("\n")은 표준 입력을 한 줄씩 읽어서 리스트로 만든다.
    # 이를 iter로 반복자로 만들고, __next__로 다음 값을 가져오는 식으로 input 함수를 재정의한다.
    # 그럼 더 빠른 input 함수를 쓸 수 있다!
    for _ in range(int(input())):
        # 테스트 케이스의 개수만큼 반복한다.
        n = int(input())
        # n을 입력받는다.
        if n == 4:
            # n이 4이면 2 2를 출력한다.
            print(2, 2)
            # 4만 따로 처리하는 이유는 이후 n을 반으로 나누어 짝수면 1을 빼서 홀수로 바꿀건데,
            # 이렇게 되면 4는 처리하지 못하기 때문이다.
        n2 = n // 2
        if n2 % 2 == 0:
            n2 -= 1
        # n을 반으로 나누어 짝수면 1을 빼서 홀수로 바꾼다.
        for prime in range(n2, 0, -2):
            # n2부터 0까지 2씩 감소하면서 반복한다.
            if che[prime] and che[n - prime]:
                # prime과 n - prime이 모두 소수이면 출력하고 다음 테스트 케이스로 넘어간다.
                print(prime, n - prime)
                break

solution()

에라토네스 체 구현 부분에 과할 정도로 공을 들인 것만 빼면 평범한 풀이이다.(모든 사람의 풀이를 보진 못해서 확신할 수는 없지만)
풀이를 잘 보면 먼저 1. 에라토네스 체를 구현하고 2. 소수를 찾아 출력하는 식으로 진행된다.
이제 여기에 함수형 패러다임을 적용해보자.

함수형 패러다임

에라토네스 체 구현

Python 에서 함수형으로 에라토네스 체를 구현하는 방법을 검색해보니 그다지 많은 자료가 나오지 않았다.
그래서 순수 함수형 언어인 Haskell 방식을 참고하여 구현해보았다.
Haskell에서 구현된 에라토네스 체는 다음과 같다.

-- 원본: https://www.literateprograms.org/sieve_of_eratosthenes__haskell_.html
<<primes_naive>>=
primes :: [Integer] -- primes는 Integer의 리스트를 반환한다.
primes = sieve [2..]  2부터 시작하는 무한 리스트를 sieve 함수에 넘긴다.
  where -- `sieve` 는 다음과 같이 정의된다.
    sieve (p:xs) = p : sieve [x|x <- xs, x `mod` p > 0]
    -- `p`를 반환하고 `xs`에서 `p`의 배수를 제거한 리스트를 다시 `sieve`에 넘긴다.

Haskell을 모른다면 굳이 이해할 필요는 없다.(밑에 파이썬으로 설명할테니까)
해당 코드는 재귀를 사용해 다음과 같이 돌아간다.

1. 2부터 시작해서 1씩 세어 나가는 반복자를 sieve 에 넘긴다.
2. 반복자를 받은 sieve 함수는 다음과 같이 작동한다.
   1. [2, 3, 4, 5, 6, ...]를 받는다.
   2. 첫번째 인자인 2를 p에 할당하고 나머지 인자인 [3, 4, 5, 6, 7...]를 xs에 할당한다.
   3. p를 반환한다.
   4. xs에서 2의 배수를 제거한 리스트 [3, 5, 7, 9, 11, ...]를 다시 sieve에 넘긴다.
3. 이를 받은 sieve 함수는 다음과 같이 작동한다.
   1. [3, 5, 7, 9, 11, ...]를 받는다.
   2. 첫번째 인자인 3을 p에 할당하고 나머지 인자인 [5, 7, 9, 11, 13, ...]를 xs에 할당한다.
   3. p를 반환한다.
   4. xs에서 3의 배수를 제거한 리스트 [5, 7, 11, 13, 15, ...]를 다시 sieve에 넘긴다.
4. 이를 받은 sieve 함수는 다음과 같이 작동한다.
   …

이제 이를 파이썬으로 구현해보자. 먼저 sieve 함수를 구현해보자.

from typing import Iterator  # Iterator 타입을 사용하기 위해 import한다.
# 뭔지 알면 따라하고 모르면 그냥 넘어가도 된다.

def sieve(xs: Iterator[int]) -> Iterator[int]:
    # `sieve` 함수는 정수 반복자를 받아서 정수 반복자를 반환한다.
    # 뭔지 모르면 `def sieve(xs):` 라고만 써도 된다.
    p = next(xs)  # `xs`의 첫번째 요소를 `p`에 할당한다.
    yield p  # p를 생성한다.
    # 위 두줄을 합쳐 `yield (p := next(xs))`로도 쓸 수 있다.
    yield from sieve(n for n in nums if n % p != 0)
    # `sieve` 함수에 `xs`에서 `p`의 배수를 제거한 반복자로부터 생성한다.

이제 이를 이용해 primes 함수를 구현하기 전에 먼저 무한 반복자인 itertools.count 함수에 대해 설명하겠다.
itertools.count(start=0, step=1)는 start부터 step씩 증가하는 무한 반복자를 반환한다.
쉽게 말해서 다음과 같은 코드와 같다.

def count(start=0, step=1):
    n = start
    while True:
        yield n
        n += step

이를 이용해 primes 함수를 구현하면 다음과 같다.

from itertools import count  # `count` 함수를 사용하기 위해 import한다.

def primes() -> Iterator[int]:
    # `primes` 함수는 정수 반복자를 반환한다.
    yield from sieve(count(2))
    # `sieve` 함수에 2부터 시작하는 반복자를 넘긴다.

# 혹은 2를 먼저 생성하고 3부터 홀수만 넘기는 방법도 있다.

def primes() -> Iterator[int]:
    yield 2
    # 2를 먼저 생성한다.
    yield from sieve(count(3, 2))
    # `sieve` 함수에 3부터 시작하는 홀수 반복자를 넘긴다.

9020번 문제에서는 최댓값이 정해져 있는 소수 리스트가 필요하기 때문에 일정 수 이하의 소수까지만 반복하는 반복자를 추가로 만들어보았다.
이를 위해 itertools.takewhile 함수를 사용했다.
itertools.takewhile(predicate, iterable)는 predicate가 True를 반환하는 동안 iterable의 요소를 반환하는 반복자를 반환한다.
다음과 같은 코드와 같다.

def takewhile(predicate, iterable):
    for x in iterable:
        if predicate(x):
            yield x
        else:
            break

예를 들어 list(takewhile(lambda x: x < 10, count(1))) == [1, 2, 3, 4, 5, 6, 7, 8, 9]이다.
여기에 primes 함수를 이용해 특정 수보다 작은 소수를 생성하는 생성자를 다음과 같이 구현할 수 있다.

from itertools import takewhile

def primes_below(n: int) -> Iterator[int]:
    yield from takewhile(lambda x: x < n, primes())
    # `lambda x: x < n` 대신 `n.__gt__`를 넘겨도 된다.
    # 또한 `primes` 함수를 생략할 수도 있다.

# 상기한 주석을 반영하면 다음과 같이도 구현할 수 있다.

def primes_below(n):
    yield 2
    yield from takewhile(n.__gt__, sieve(count(3, 2)))

구현한 함수들을 이용해 소수 판별 함수를 다음과 같이 구현할 수 있다.

from typing import Iterator
from itertools import count, takewhile

def sieve(nums: Iterator[int]) -> Iterator[int]:
    n = next(nums)
    yield n
    yield from sieve(i for i in nums if i % n != 0)


def primes() -> Iterator[int]:
    yield 2
    yield from sieve(count(3, 2))


def primes_below(n: int) -> Iterator[int]:
    yield from takewhile(n.__gt__, primes())


is_prime = set(primes_below(10000)).__contains__

본문

이제 본론으로 들어가보자.
원래 코드는 다음과 같다.

for _ in range(int(input())):
    n = int(input())
    if n == 4:
        print(2, 2)
    n2 = n // 2
    if n2 % 2 == 0:
        n2 -= 1
    for prime in range(n2, 0, -2):
        if che[prime] and che[n - prime]:
            print(prime, n - prime)
            break

먼저 n이 4가 아닌 경우부터 고려하자.
먼저 n을 반으로 나눈 n2를 구하고, n2가 짝수라면 1을 빼줘야한다.
하지만 사실 if 문 쓸 것도 없이 n2 = n // 2 - 1 + (n // 2) % 2로 바로 구할 수 있다.
이제 n2부터 2씩 감소하면서 prime이 소수이고 n - n2도 소수인 n2를 찾으면 된다.
zip을 이용해 둘을 짝짓고, filter를 이용해 둘다 소수인 것만 걸러내고, next를 이용해 첫 번째 값을 가져오면 된다.
먼저 둘다 소수인지 판정하는 함수는 is_prime을 이용해 구현할 수 있다.

def is_both_prime(nm):
    return all(map(is_primes, nm))

이제 zip과 filter를 이용해 n2를 구할 수 있다.

p, q = next(filter(is_both_prime, zip(range(np2, 0, -2), range(n - np2, n, 2))))

# 너무 기니까 변수를 분리하면

np2_and_n_m_np2 = zip(range(np2, 0, -2), range(n - np2, n, 2))
p, q = next(filter(is_both_prime, np2_and_n_m_np2))

이 과정을 함수로 만들면 다음과 같다.

def is_both_prime(nm):
    return all(map(is_primes, nm))


def divide_not_4(n):
    np2 = n // 2 - 1 + (n // 2) % 2
    np2_and_n_m_np2 = zip(range(np2, 0, -2), range(n - np2, n, 2))
    p, q = next(filter(is_both_prime, np2_and_n_m_np2))
    return f"{p} {q}"
    # 문자열로 출력해야하므로 f-string을 이용해 출력한다.

이제 n이 4인 경우 또한 고려할 수 있는 함수를 만들면 된다.

def divide(n):
    return "2 2" if n == 4 else divide_not_4(n)

이제 입력을 받아 출력하는 함수를 만들면 된다.

def solution():
    next(input := map(int, open(0).read().split()))
    # input을 받아서 정수로 바꾸어준다.
    # 첫 입력인 테스트 케이스의 개수는 필요없으므로 버린다.
    print(*map(divide, input), sep="\n")
    # 테스트 케이스 별 답을 출력한다.

풀이

최종적으로 제출한 코드는 다음과 같다.

import sys
from itertools import count, takewhile

sys.setrecursionlimit(10000)
# 재귀 한도를 늘려준다.

def sieve(nums):
    yield (p := next(nums))
    yield from sieve(n for n in nums if n % p != 0)


def primes_below(n):
    # `primes`를 생략하고 바로 `primes_below`를 정의했다.
    yield 2
    yield from takewhile(n.__gt__, sieve(count(3, 2)))


is_primes = set(primes_below(10000)).__contains__


def is_both_prime(nm):
    return all(map(is_primes, nm))


def divide_not_4(n):
    np2 = n // 2 - 1 + (n // 2) % 2
    np2_and_n_m_np2 = zip(range(np2, 0, -2), range(n - np2, n, 2))
    p, q = next(filter(is_both_prime, np2_and_n_m_np2))
    return f"{p} {q}"


def divide(n):
    return "2 2" if n == 4 else divide_not_4(n)


def solution():
    next(input := map(int, sys.stdin.read().split()))
    print(*map(divide, input), sep="\n")


solution()

결과는 200ms로 기존 코드의 결과인 52ms보다 느리다.
sieve 함수는 꼬리재귀를 사용하는데, 파이썬에서는 꼬리재귀를 최적화하지 않기 때문에 영향을 미친 것 같다.
하지만 재밌는 과정이었다!
이미 풀어본 문제도 함수형 패러다임을 적용해봐야겠다.

리스트 컴프리헨션과 filter(221219 추가)

8번째 줄의 n for n in nums if n % p != 0를 filter(p.__rmod__, nums)로 변경해보았다.
해당 풀이는 172ms 라는 무려 28ms나 단축된 결과를 얻을 수 있었다!
리스트 컴프리헨션이 filter보다 읽기 더 편한건 사실이지만 아직 성능적으로는 아쉬운 부분이 있는 것 같다.
더 많은 코드를 작성해보며 적재적소에 맞춰 사용하자!

set.issuperset (221219 추가)

is_both_prime을 따로 만들지 않고 set.issuperset 메소드를 적용해보았다.

are_primes = set(primes_below(10000)).issuperset


def divide_not_4(n):
    np2 = n // 2 - 1 + (n // 2) % 2
    np2_and_n_m_np2 = zip(range(np2, 0, -2), range(n - np2, n, 2))
    p, q = next(filter(are_primes, np2_and_n_m_np2))
    return f"{p} {q}"

해당 풀이는 168ms로 4ms 단축된 결과를 얻었다.