[Python3] a^3 + b^3 = c^3 + d^3 을 만족하는 자연수 1000 이하의 조합

맥도웰 책을 샀는데 33페이지에 잠깐 언급이 있어서 풀어봤다. 생각보다 만족하는 조합이 엄청 많이 나와서 놀랐다.

#!/usr/bin/env python3

from itertools import combinations as comb

cube = [x**3 for x in range(1001)]

dp1 = dict()

for a in range(1, 1000):

    for b in range(a+1, 1001):

        dp1[(a, b)] = cube[a] + cube[b]

dp2 = dict()

for k, v in dp1.items():

    dp2.setdefault(v, []).append(k)

res = []

for k, v in dp2.items():

    if len(v) == 2:

        res.append(v)

    elif len(v) > 2:

        for com in comb(v, 2):

            res.append(list(com))

for item in res:

    a, b = item[0]

    c, d = item[1]

    text = '{}**3 + {}**3 == {}**3 + {}**3 == {}'.format(

        a, b, c, d, dp1[item[0]])

    print(text)

조금 더 생각해보니 위 코드에 비효율적인 부분이 있다.

첫째는 dict를 만들어놓고 k,v를 뒤집기 위해 dict를 다시 만든 것이고 (처음부터 뒤집으면 되는데)

둘째는 출력 전에 res라는 list를 만든 것이고 (바로 출력해도 되는데)

셋째는 itertools의 combination을 import 했다는 것이다. (그냥 for문 사용하면 되는데)

아래는 세가지 비효율을 제거한 버전이다.

#!/usr/bin/env python3

cube = [x**3 for x in range(1001)]

dp = dict()

for a in range(1, 1000):

    for b in range(a+1, 1001):

        dp.setdefault(cube[a] + cube[b], []).append((a, b))

for k, v in dp.items():

    if len(v) > 1:

        for i in range(len(v)-1):

            for j in range(i+1, len(v)):

                print('{}**3 + {}**3 == {}**3 + {}**3 == {}'.format(

                    *v[i], *v[j], k))

조금 더 생각해보니 a**3 + b**3 연산 결과를 미리 dictionary에 넣어둬봐야 별 실익이 없어보인다. 그리고 len(v) > 1 부분도 굳이 확인하지 않아도 되겠다. 어차피 len(v) == 1 이면 for j in range(i+1, len(v)) 부분이 []이 되어 안으로 진입하지 않을테니까...

#!/usr/bin/env python3

cubic = {}

for a in range(1, 1000):

    for b in range(a+1, 1001):

        cubic.setdefault(a**3 + b**3, []).append((a, b))

output_str = '{}**3 + {}**3 == {}**3 + {}**3 == {}'

for k, v in cubic.items():

    for i in range(len(v)-1):

        for j in range(i+1, len(v)):

            print(output_str.format(*v[i], *v[j], k))