프로젝트 오일러 049

세 항이 소수이면서 다른 수의 순열이 되는 4자리 수의 등차수열

프로젝트 오일러 049
Photo by ALINA MATVEYCHEVA / Unsplash

문제

49번 문제
세 항이 소수이면서 다른 수의 순열이 되는 4자리 수의 등차수열 찾기

4자리 소수를 모두 찾는 것은 소수체를 만들면 아주 빠르게 찾을 수 있습니다. 같은 숫자로 이루어진 순열이면서 등차수열을 이루는 4자리 소수들을 찾기 위해서는 여러 가지 방법이 있을 수 있습니다.

4자리 소수들을 같은 순열끼리 묶기

같은 숫자들의 순열로 되어 있는 수들은 각 자리 수를 정렬하면 같은 결과가 됩니다. 이 정렬결과들을 키로하고, 순열들의 목록을 값으로 하는 사전을 만든다면, 같은 순열끼리 소수를 묶을 수 있습니다. 이후 각 묶음 내에서 등차수열을 이루는 3개의 값이 있는지를 검사하여 답을 찾을 수 있습니다.

일련의 수열에서 등차수열을 이루는 3개의 값을 찾을 때에는 크기순으로 정렬된 a, b, c 세 수에서 b × 2 = a + c 인지를 검사하면 됩니다.

소수체 만드는 함수와 조합을 만드는 제너레이터는 이전 문제 풀이에서 작성한 바 있으니, 이들을 utils라는 모듈로 따로 정리해놓고 사용한다고 가정합니다.

from utils import seive, combinations


def helper(n: int) -> tuple[int, ...]:
    return tuple(map(int, sorted(str(n))))


def main():
    s = filter(lambda x: x > 999, seive(9999))
    d = {}
    # 각 4자리 소수를 순열별로 정리
    for p in s:
        k = helper(p)
        d.setdefault(k, []).append(p)
    # 각 그룹에서 등차수열 찾기
    for k in d:
        # 3개 미만의 그룹은 통과
        if len(d[k]) < 3:
            continue
        # 문제의 예시 케이스는 통과
        if d[k][0] == 1487:
            continue
        for (a, b, c) in combinations(d[k], 3):
            if b * 2 == a + c:
                print("".join(map(str, (a, b, c))))
                return


if __name__ == "__main__":
    main()