삼각수, 사각수, 오각수 같은 다각수들은 아래의 공식으로 만들 수 있습니다.
- 삼각수 – P3n = n * ( n + 1) / 2 >> 1, 3, 5, 10, 15, …
- 사각수 – P4n = n * n >> 1, 4, 9, 16, 25, …
- 오각수 – P5n = (n * ( n * 3 – 1)) / 2 > 1, 5, 12, 22, 35, …
- 육각수 – P6n = (n * (n * 2 – 1)) > 1, 6, 15, 28, 45, …
- 칠각수 – P7n = ((n * (n * 5 – 3)) / 2 > 1, 7, 18, 34, 55, …
- 팔각수 – P8n = n * (n * 3 – 2) > 1, 8, 21, 40, 65, …
그런데 4자리 숫자 8128, 2882, 8281 (순서대로 각각 3, 5, 4각수)에는 재미있는 성질이 있습니다. 먼저 각 숫자들은 두 자리씩 꼬리를 물고 진행합니다. 그리고 각 숫자들은 모두 서로 다른 다각수입니다. 이런 성질을 갖는 네자리 숫자 세 개는 이들이 유일합니다. 이렇게 순환하면서 서로 다른 다각수가 되는 4자리 숫자 여섯개의 유일한 순서쌍을 찾고 그 합을 구하세요.
두 네자리 수 a, b 가 꼬리를 문다는 것은 a를 100으로 나눈 나머지와, b를 100으로 나눈 몫이 서로 같다는 것이다. 네 자리 다각수 전체의 집합이 있다면, 이 조건을 사용해서 각각의 다각수에서 꼬리를 물 수 있는 다각수들을 연결하는 맵을 작성하는 일은 간단하다. 그리고 네 자리 다각수를 전부 구해도 그 수가 그리 많지 않을 것이기 때문에 빠르게 끝날 것으로 예상된다. 연결을 만들 수 있는 맵만 만든다면 그 이후의 일은 간단하다. 임의의 한 수로부터 연결 가능한 다각수 5개를 더 찾은 다음, 마지막 다각수로부터 첫번째 다각수가 다시 연결될 수 있는지를 검사하면 된다. 문제에서는 순환고리를 만들 수 있는 다각수의 순서쌍은 유일하다고 했으니, 첫번째로 조건을 만족하는 집합을 찾으면 된다.
중간에 더 이상 다른 다각수를 찾을 수 없다면 거슬러 올라가서 다음 수로 시도해보아야 한다. 이러한 흐름은 백트래킹으로 구현하는데, 파이썬에서 백트래킹을 구현하는 가장 간단한 방법은 yield from
구문을 사용해서 제너레이터의 위임 구조를 재귀적으로 만드는 것이다. 이 블로그에서 순열, 조합을 제너레이터를 사용하여 구현하는 방법을 소개한 적이 있는데, 그 글에서도 같은 방식을 사용했고 코드가 놀라울 수준으로 간단해지는 것을 확인할 수 있다.
전체 다각수 집합 및 연결 맵 만들기
사실 모든 육각수는 삼각수이기도 하므로, 각각의 다각수에는 이 수가 어떤 n각수인지를 구분하는 정보가 필요하다. 이 두 가지 정보를 튜플로 묶어서 생성 가능한 전체 다각수를 만들어 보자.
from typing import Generator
def gen() -> Generator[tuple[int, int], None, None]:
equations = [
lambda n: (n * (n + 1)) // 2,
lambda n: n ** 2,
lambda n: (n * (3 * n - 1)) // 2,
lambda n: n * (2 * n - 1),
lambda n: (n * (5 * n - 3)) // 2,
lambda n: n * (n * 3 - 2)
]
for i, eq in enumerate(equations):
k = 1
while True:
x = eq(k)
if 999 < x <= 9999:
yield (i + 3, x)
elif x > 9999:
break
k += 1
numbers = list(gen())
전체 다각수 집합을 만들었으니, 이제 각 다각수에 대해서 연결 가능한 다각수의 집합을 구성해보자. 이것도 간단하다. 아래 예와 달리 for 루프를 돌면서 리스트에 append 할 수도 있겠지만, 이게 더 빠를 것이다.
links = {x:[y for y in numbers if x[0] != y[0] and x[1] % 100 == y[1] // 100]
for x in numbers}
백트래킹으로 순서쌍 찾기
이제 연결되는 순서쌍을 구하는 핵심 로직을 구현한다.
def helper(
d: dict[tuple[int, int], list[tuple[int, int]]], # 링크 맵
acc: list[tuple[int, int]], # 현단계까지 모인 순서쌍
n: tuple[int, int], # 이번 단계의 값
k: int # 재귀 깊이
) -> list[tuple[int, int]]:
if k == 1:
# n 이 여섯번째 다각수인 경우 첫번째 다각수와 연결되는지 확인
if acc[0][1] % 100 == n[1] // 100:
yield [*acc, n]
for m in d.get(n, []):
# m이 이미 추가된 다각수 유형인지 확인
if m[0] in [a[0] for a in acc]:
continue
yield from helper(d, [*acc, n], m, k-1)
yield from
을 사용하면 리턴해줄 값이 없는 경우는 자동으로 무시되기 때문에 백트래킹 로직을 구현하는 것이 매우 간편해짐을 알 수 있다. 참고로 이중 for 루프를 탈출하려 할 때 아래 테크닉을 참고해두면 좋다. 내부의 루프가 break
로 중단되는 경우 else:
절이 실행되지 않기 때문에, else:
에 continue
를 쓰고 바로 뒤에서 브레이크를 걸면 내부 루프가 break로 중단되면 바깥쪽 루프도 같이 중단된다.
for num in numbers:
for xs in helper(links, [], num, 6):
print(sum(x[1] for x in xs))
break
else:
continue
break