[코드 트리] 꼬리잡기놀이 - python

문제 링크

https://www.codetree.ai/frequent-problems/tail-catch-play/description

코드트리

풀이 

해설을 참고하여 풀었다. 이런 문제가 나왔을 때, 고려해야 할 점들이 많기 때문에 아이디어를 쉽게 생각하지 못하는 것 같다. 

 

rail이라는 배열에는 index(팀 번호) 별 rail의 위치를 저장하였다. 이때, rail[0]은 머리의 위치이고 머리를 시작으로 2번의 위치 -> 3번의 위치 -> 4번의 위치를 대입해주었다. 나중에 한칸씩 이동할 때 편리하게 하기 위해서 deque를 사용했다. 팀 번호는 graph를 탐색하면서 머리가 나오는 순서대로 번호를 매겨준 것이다. 

tail이라는 배열에는 위에서 구한 rail에서 꼬리 위치의 인덱스 번호를 넣어주었다. 즉, 머리에서부터 몇번째인지 

board 라는 배열에는 각 rail에 해당하는 위치들을 팀 번호로 표시해주었다. 

 

1. find_rail

각 팀 별 rail의 인덱스 값을 저장하고, board에 팀 번호를 표시해준다. tail에 팀 별 꼬리의 인덱스 번호도 넣어준다.

처음에 graph를 입력받을 때, 1의 위치(머리의 위치)를 팀 번호 별로 넣어주었다.

rail은 겹치는 경우가 없기 때문에 머리를 기준으로 2, 3, 4, 순서로 찾아주면 된다.

def find_rail(num):
	
    # rail 별 머리 위치 
    s_x, s_y = rail[num][0]
    
    # 탐색을 위한 방문 표시 
    visited[s_x][s_y] = True
	# board에 레일 번호로 레일 표시 
	board[s_x][s_y] = num
    
    # 시작 위치 대입 
    q = deque()
    q.append((s_x,s_y))
    # tail 인덱스 번호를 기록하기 위함 
    tmp_idx = 1

    while q:
        x, y = q.popleft()

        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]

            if nx <= 0 or nx > n or ny <= 0 or ny > n:
                continue
			
            # 2번 찾기 
            if not visited[nx][ny] and graph[nx][ny] == 2:
                visited[nx][ny] = True
                # rail에 넘어줌 
                rail[num].append((nx, ny))
                # board에 표시 
                board[nx][ny] = num
                q.append((nx,ny))
                tmp_idx += 1
                s_x = nx
                s_y = ny
	
    # 3번(꼬리) 찾기 
    for i in range(4):
        snx = s_x + dx[i]
        sny = s_y + dy[i]

        if snx <= 0 or snx > n or sny <= 0 or sny > n:
            continue

        if graph[snx][sny] == 3 :
            rail[num].append((snx, sny))
            s_x = snx
            s_y = sny
            visited[s_x][s_y] = True
            board[s_x][s_y] = num
            # 꼬리가 머리에서부터 몇번째인지 
            tail[num] = tmp_idx
            break

    q = deque()
    q.append((s_x,s_y))
	
    # 4번, 즉 레일 계속 찾기 
    while q:
        x, y = q.popleft()

        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]

            if nx <= 0 or nx > n or ny <= 0 or ny > n:
                continue

            if not visited[nx][ny] and graph[nx][ny] == 4 :
                visited[nx][ny] = True
                # rail에 추가 
                rail[num].append((nx,ny))
                # board에 번호 표시 
                board[nx][ny] = num
                q.append((nx,ny))

이제 k 번동안 아래 과정을 반복한다. 

1. 각 팀 별 머리 사람을 따라 이동한다. 

1) rotate를 통해 rail를 한 칸씩 이동시켜준다. 

    for i in range(1,m+1):
        # 처음 인덱스가 머리 위치
        rail[i].rotate(1)

2) graph에도 이동한 내용을 반영시켜준다. 

rail[팀번호][0]에는 머리 사람의 위치가 저장되어있고, tail에 저장된 꼬리의 인덱스보다 작은 곳에는 머리사람과 꼬리사람이 아닌 나머지 사람(2번 사람), tail[팀번호]에는 꼬리 사람의 인덱스 번호(머리사람으로부터 몇번째 위치)가 저장되어있으므로 아래와 같이 옮겨주면 된다. 

    for i in range(1,m+1):

        for j in range(len(rail[i])):
            if j == 0 :
                graph[rail[i][j][0]][rail[i][j][1]] = 1
            elif j < tail[i]  :
                graph[rail[i][j][0]][rail[i][j][1]] = 2
            elif j == tail[i] :
                graph[rail[i][j][0]][rail[i][j][1]] = 3
            else :
                graph[rail[i][j][0]][rail[i][j][1]] = 4

3. 공을 던져 점수를 얻는다.

현재 turn이 몇번째 turn인지에 따라 방향과 기준 행이나 열을 정하고, 사람이 있는지 검사해준다. 

사람이 있다면(1~3 값을 만난다면) 점수를 계산해주고, 해당 팀의 머리사람과 꼬리 사람을 변경한다. 

def throw():
    global turn

    t = (turn-1) % (4*n) + 1
	
    # 오른쪽으로 던짐 
    if t <= n :
        for i in range(1,n+1):
       		# 사람을 만나면 
            if 1 <= graph[t][i] <= 3 :
                # 얻는 점수 계산 함수 
                get_point(t,i)
                # 머리 사람과 꼬리 사람을 변경하는 함수 
                reverse(board[t][i])
                break
	# 위쪽으로 던짐 
    elif t <= 2 * n :
        t -= n
        for i in range(1,n+1):
            if 1 <= graph[n+1-i][t] <= 3 :
                get_point(n+1-i, t)
                reverse(board[n+1-i][t])
                break
	# 왼쪽으로 
    elif t <= 3 * n:
        t -= (2*n)
        for i in range(1,n+1):
            if 1 <= graph[n+1-t][n+1-i] <= 3:
                get_point(n+1-t, n+1-i)
                reverse(board[n+1-t][n+1-i])
                break
    # 아래쪽으로 던짐
    else :
        t -= (3*n)
        for i in range(1,n+1):
            if 1 <= graph[i][n+1-t] <= 3:
                get_point(i, n+1-t)
                reverse(board[i][n+1-t])
                break

얻는 점수를 계산하는 함수는 다음과 같다. 

def get_point(x,y):
    global ans

    # 팀 번호
    num = board[x][y]
    # 팀에서 몇번째 사람인지
    cnt = rail[num].index((x,y))
    ans += (cnt +1) * (cnt+1)

머리사람과 꼬리사람을 바꾸는 함수는 다음과 같다.

def reverse(num):
	
    new = deque()

    # 꼬리부터 머리 사람 대입 -> 이젠 꼬리 사람이 머리, 머리 사람이 꼬리임 
    for i in range(tail[num], -1, -1):
        new.append(rail[num][i])

    # 남은 거 반대로 대입 
    for i in range(len(rail[num])-1, tail[num], -1):
        new.append(rail[num][i])

    rail[num] = new
	
    # graph에도 반영해줌 
    for i in range(len(rail[num])):
        if i == 0 :
            graph[rail[num][i][0]][rail[num][i][1]] = 1
        elif i < tail[num] :
            graph[rail[num][i][0]][rail[num][i][1]] = 2
        elif i == tail[num] :
            graph[rail[num][i][0]][rail[num][i][1]] = 3
        else :
            graph[rail[num][i][0]][rail[num][i][1]] = 4

4. turn + 1 을 한다. 

 

전체 코드는 다음과 같다. 

from collections import deque

dx = [0,-1,0,1]
dy = [1,0,-1,0]

# 팀별 레일 정보 저장 함수 
def find_rail(num):

    s_x, s_y = rail[num][0]
    visited[s_x][s_y] = True
    board[s_x][s_y] = num
    q = deque()
    q.append((s_x,s_y))
    tmp_idx = 1

    while q:
        x, y = q.popleft()

        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]

            if nx <= 0 or nx > n or ny <= 0 or ny > n:
                continue

            if not visited[nx][ny] and graph[nx][ny] == 2:
                visited[nx][ny] = True
                rail[num].append((nx, ny))
                board[nx][ny] = num
                q.append((nx,ny))
                tmp_idx += 1
                s_x = nx
                s_y = ny

    for i in range(4):
        snx = s_x + dx[i]
        sny = s_y + dy[i]

        if snx <= 0 or snx > n or sny <= 0 or sny > n:
            continue

        if graph[snx][sny] == 3 :
            rail[num].append((snx, sny))
            s_x = snx
            s_y = sny
            visited[s_x][s_y] = True
            board[s_x][s_y] = num
            tail[num] = tmp_idx
            break

    q = deque()
    q.append((s_x,s_y))

    while q:
        x, y = q.popleft()

        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]

            if nx <= 0 or nx > n or ny <= 0 or ny > n:
                continue

            if not visited[nx][ny] and graph[nx][ny] == 4 :
                visited[nx][ny] = True
                rail[num].append((nx,ny))
                board[nx][ny] = num
                q.append((nx,ny))
                
# 팀 별 머리 사람을 따라 한 칸씩 이동하는 함수 
def move():

    for i in range(1,m+1):
        # 처음 인덱스가 머리 위치 -> 한칸씩 이동 
        rail[i].rotate(1)
	
    # graph에도 반영 
    for i in range(1,m+1):

        for j in range(len(rail[i])):
            if j == 0 :
                graph[rail[i][j][0]][rail[i][j][1]] = 1
            elif j < tail[i]  :
                graph[rail[i][j][0]][rail[i][j][1]] = 2
            elif j == tail[i] :
                graph[rail[i][j][0]][rail[i][j][1]] = 3
            else :
                graph[rail[i][j][0]][rail[i][j][1]] = 4

# 머리 사람과 꼬리 사람을 바꾸는 함수 
def reverse(num):

    new = deque()

    # 꼬리부터 머리 사람
    for i in range(tail[num], -1, -1):
        new.append(rail[num][i])

    # 남은 거 반대로 대입 
    for i in range(len(rail[num])-1, tail[num], -1):
        new.append(rail[num][i])

    rail[num] = new

	# graph에도 반영 
    for i in range(len(rail[num])):
        if i == 0 :
            graph[rail[num][i][0]][rail[num][i][1]] = 1
        elif i < tail[num] :
            graph[rail[num][i][0]][rail[num][i][1]] = 2
        elif i == tail[num] :
            graph[rail[num][i][0]][rail[num][i][1]] = 3
        else :
            graph[rail[num][i][0]][rail[num][i][1]] = 4

# 공을 던져서 얻는 점수 계산 함수 
def get_point(x,y):
    global ans

    #팀 번호
    num = board[x][y]
    # 몇번째 사람인지
    cnt = rail[num].index((x,y))
    ans += (cnt +1) * (cnt+1)

# 공 던지는 함수 
def throw():
    global turn
	
    # 현재 몇번째 turn 인지 
    t = (turn-1) % (4*n) + 1

    if t <= n :
        for i in range(1,n+1):
            if 1 <= graph[t][i] <= 3 :
                get_point(t,i)
                reverse(board[t][i])
                break

    elif t <= 2 * n :
        t -= n
        for i in range(1,n+1):
            if 1 <= graph[n+1-i][t] <= 3 :
                get_point(n+1-i, t)
                reverse(board[n+1-i][t])
                break

    elif t <= 3 * n:
        t -= (2*n)
        for i in range(1,n+1):
            if 1 <= graph[n+1-t][n+1-i] <= 3:
                get_point(n+1-t, n+1-i)
                reverse(board[n+1-t][n+1-i])
                break
    else :
        t -= (3*n)
        for i in range(1,n+1):
            if 1 <= graph[i][n+1-t] <= 3:
                get_point(i, n+1-t)
                reverse(board[i][n+1-t])
                break

n, m, k = map(int,input().split())

graph = [ [0] * (n+1) ]

rail = [  deque() for _ in range(m+1)]
tail = [0] * (m+1)
idx = 1
visited = [ [False] * (n+1) for _ in range(n+1) ]
board = [ [0] * (n+1) for _ in range(n+1)]
turn = 1
ans = 0

for i in range(1,n+1):
    graph.append([0] + list(map(int,input().split())))
    for j in range(1,n+1):
        if graph[i][j] == 1:
            rail[idx].append((i,j))
            idx += 1

for i in range(1,m+1):
    find_rail(i)

for _ in range(k):
    move()

    throw()
    turn += 1

print(ans)