문제

삼각수 T_n(n ≥ 1)는 [그림]에서와 같이 기하학적으로 일정한 모양의 규칙을 갖는 점들의 모음으로 표현될 수 있다.

[그림]

자연수 n에 대해 n ≥ 1의 삼각수T_n는 명백한 공식이 있다.

T_n = 1 + 2 + 3 + … + n = n(n+1)/2

1796년, 가우스는 모든 자연수가 최대 3개의 삼각수의 합으로 표현될 수 있다고 증명하였다. 예를 들어,

• 4 = T₁ + T₂
• 5 = T₁ + T₁ + T₂
• 6 = T₂ + T₂ or T₆ = T₃
• 10 = T₁ + T₂ + T₃ or T₁₀ = T₄

이 결과는 증명을 기념하기 위해 그의 다이어리에 “Eureka! num = Δ + Δ + Δ” 라고 적은 것에서 유레카 이론으로 알려졌다. 꿍은 몇몇 자연수가 정확히 3개의 삼각수의 합으로 표현될 수 있는지 궁금해졌다. 위의 예시에서, 5와 10은 정확히 3개의 삼각수의 합으로 표현될 수 있지만 4와 6은 그렇지 않다.

자연수가 주어졌을 때, 그 정수가 정확히 3개의 삼각수의 합으로 표현될 수 있는지 없는지를 판단해주는 프로그램을 만들어라. 단, 3개의 삼각수가 모두 달라야 할 필요는 없다.

입력

프로그램은 표준입력을 사용한다. 테스트케이스의 개수는 입력의 첫 번째 줄에 주어진다. 각 테스트케이스는 한 줄에 자연수 K (3 ≤ K ≤ 1,000)가 하나씩 포함되어있는 T개의 라인으로 구성되어 있다.

출력

프로그램은 표준출력을 사용한다. 각 테스트케이스에대해 정확히 한 라인을 출력한다. 만약 K가 정확히 3개의 삼각수의 합으로 표현될수 있다면 1을, 그렇지 않다면 0을 출력한다.

삼각수는 곧 해당 숫자까지 모든 숫자의 합이다. 이는 (n-1) * n // 2란 식으로 수렴된다. 삼각수를 구하는 건 어렵지 않다. 문제는, 어떤 숫자가 3개의 삼각수의 합으로 이뤄져 있는지 여부를 판단하는 것.

어떻게 풀어야 할지 막막했는데, 알고리즘 분류를 보니 브루트포스다. 우직하게 처음부터 끝까지 돌며 찾아내라는 얘기.

1부터 1000까지(정확히는 1000이 막 넘는 삼각수까지) 삼각수를 리스트로 만들어두고 3중 반복문을 돌린다. 정답(ans)의 초기값은 0. 세 삼각수의 합이 입력한 숫자와 일치하면 정답을 ‘1’로 바꾸고 반복문을 종료. 마지막으로 정답을 출력해주면 끝.

시간초과를 피할 수 있을까. 아슬아슬하게 통과했다.

w = [i*(i+1)//2 for i in range(1, 46)] 
for _ in range(int(input())):
    n = int(input())
    r = int((n*2)**0.5)+1 # 입력한 숫자까지만 반복문을 돌리면 된다.
    ans = 0
    for i in range(r):
        for j in range(r):
            for k in range(r):
                if n == w[i] + w[j] + w[k]:
                    ans = 1
                    break
    print(ans)

[문제 보기]