백준 #9461 - 파도반수열 python

< 문제 >

오른쪽 그림과 같이 삼각형이 나선 모양으로 놓여져 있다. 첫 삼각형은 정삼각형으로 변의 길이는 1이다.

그 다음에는 다음과 같은 과정으로 정삼각형을 계속 추가한다.

나선에서 가장 긴 변의 길이를 k라고 했을 때, 그 변의 길이가 k인 정삼각형을 추가한다.

파도반 수열 P(N)은 나선에 있는 정삼각형의 변의 길이다. 

P(1)부터 P(10)까지 첫 10개의 숫자는 1,1,1,2,2,3,4,5,7,9 이다.

< 입력 > 

test case T 주어지고

N이 주어진다. (1<= N <= 100)

P(N)을 출력하라.

def rule(n,s):
    if n <= 5 or s[n]!=0:
        return s[n]
    else :
        s[n]= rule(n-1,s)+rule(n-5,s)
        return s[n]
 
if __name__ == '__main__':
    T = int(input())
    for _ in range(T):
        n = int(input())
        if n<=5:
            s=[0 for _ in range(6) if True]
        else:
            s=[0 for _ in range(n+1) if True]
        s[1] = s[2] = s[3] = 1
        s[4] = s[5] = 2
 
        print(rule(n,s))