[jungol] 5398 : 달팽이트리- python code

문제

달팽이의 나라 달팽국에서는 크리스마스 때 특별한 트리를 세워 장식한다.
이 트리를 달팽국 사람들은 달팽이 트리라고 부르는데,
달팽이 트리는 꼭대기부터 0~9까지의 정수가 반시계 방향으로 밖에서 안으로 채워진 형태를 띄고 있다.
달팽이 트리의 크기는 매 해 달팽국의 행복도에 비례하여 다르게 정해진다.

예를 들어 아래와 같이 행복도가 6이면 크기 6짜리 달팽이 트리가 만들어진다.

올해의 달팽국 행복도가 주어지면 그에 알맞는 크기의 달팽이 트리를 출력하는 프로그램을 작성하자.

입력형식
삼각형의 한 변의 길이 N(N의 범위는 100 이하의 양의 정수)을 입력받는다.

출력형식
주어진 형태대로 높이가 N인 달팽이 삼각형을 출력한다. 숫자 사이는 한 개의 공백으로 구분한다.

처음에는 이런식으로 코드를 짜보았다.

n = int(input())
table = [ [0 for x in range(2*n-1)] for y in range(n)]
X,Y,num,dir = n-1,0,1,0
dx=[-1,2,-1]
dy=[1,0,-1]

while True:
    table[Y][X] = num
    t_X, t_Y = X+dx[dir], Y+dy[dir]
    if 0<=t_X<2*n and 0<= t_Y <n and table[t_Y][t_X] == 0:
        X,Y = t_X,t_Y
    else:
        dir = (dir+1)%3
        X,Y = X+dx[dir], Y+dy[dir]
        if table[Y][X] != 0:
            break
    num += 1

for y in range(n):
    for x in range(2*n-1):
        #print(table[y][x],end='')
        if table[y][x] != 0:
            print((table[y][x]-1)%10, end='')
        else:
            print(' ', end='')
    print()

그런데 여기서는 문제가 있었다.
바로 n=1 일때는 IndexError가 뜬다는 것이다.
그 이유는 만약 n=1이라면 dir이 어떻든 어떻게 움직여도 이미 인덱스가 범위를 벗어났기 때문이다.

그래서 while문의 중단조건에 다른 조건을 추가 시켜주었다.

n = int(input())
table = [ [0 for x in range(2*n-1)] for y in range(n)]
X,Y,num,dir = n-1,0,1,0
dx=[-1,2,-1]
dy=[1,0,-1]

while True:
    table[Y][X] = num
    t_X, t_Y = X+dx[dir], Y+dy[dir]
    if 0<=t_X<2*n and 0<= t_Y <n and table[t_Y][t_X] == 0:
        X,Y = t_X,t_Y
    else:
        dir = (dir+1)%3
        X,Y = X+dx[dir], Y+dy[dir]
        if not 0<=X<2*n or not 0<= Y <n or table[Y][X] != 0 :
            break
    num += 1

for y in range(n):
    for x in range(2*n-1):
        #print(table[y][x],end='')
        if table[y][x] != 0:
            print((table[y][x]-1)%10, end='')
        else:
            print(' ', end='')
    print()

그러면 n=1인 경우도 해결되었다.

완성! :)