[코드트리] 루돌프의 반란 - Python

문제링크

https://www.codetree.ai/training-field/frequent-problems/problems/rudolph-rebellion/description?page=1&pageSize=20 

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

풀이

구현 문제로, 산타가 있는 위치로 밀려났을 때 연쇄적으로 밀려다도록 하는 부분은 재귀 함수 사용해서 풀었다. 

 

필요한 변수들은 다음과 같다. 

n, m, p, c, d = map(int,input().split())
r_x, r_y = map(int,input().split())
r_x -=1 
r_y -=1

# index에 해당하는 산타의 위치 저장 
santa = [ [] for _ in range(p+1) ] 
# 위치별 산타 번호 저장 
graph = [ [0] * n for _ in range(n) ]
# index에 해당하는 산타의 점수 저장 
scores = [0] * (p+1)
# index에 해당하는 산타 존재 여부 저장 
out = [False] * (p+1)
# index에 해당하는 산타가 언제까지 기절하는지 시간 저장 
stop_time = [0] * (p+1)
# 몇 초가 흘렸는지 
t = 1 

for _ in range(p):
    i, a, b = map(int,input().split())
    a -= 1 
    b -= 1 
    santa[i] = [a,b]
    graph[a][b] = i

연쇄적으로 밀어내는 함수는 다음과 같다.

# 산타 번호, 산타가 이동할 위치, 이동한 방향 
def shoot(s_id, x, y, direct):
    # 밀어내진 위치가 범위를 벗어나면 out 
    if x < 0 or x >= n or y < 0 or y >= n:
        out[s_id] = True 
        return 
    
    # 밀어난 위치에도 산타가 있으면 재귀적으로 밀어낸다 
    if graph[x][y] :
        tmp_id = graph[x][y]
        # 착지 
        graph[x][y] = s_id
        santa[s_id] = [x,y]
        # 연쇄적으로 밀어냄 
        shoot(tmp_id, x+dx[direct], y+dy[direct], direct)
    # 밀어난 위치가 빈 칸이라면 바로 착지 
    else :
        graph[x][y] = s_id
        santa[s_id] = [x,y]

 

1. 루돌프 이동 함수

# 상하좌우대각선 8방향 
dx = [-1,0,1,0,1,1,-1,-1]
dy = [0,1,0,-1,1,-1,1,-1]

def move_rudolf():
    global r_x, r_y 

    # 산타 별 거리 구하기
    route = []
    for i in range(1,p+1):
        if out[i]:
            continue 

        s_x, s_y = santa[i]
        dist = (s_x-r_x)**2 + (s_y-r_y)**2 
        # 거리 작은 순 -> r 큰 순 -> c 큰 순으로 정렬 
        route.append([dist, -s_x, -s_y, i])

    route.sort()
    
    # 거리, 돌진할 산타 위치, 돌진한 산타 번호
    dist, x, y, i = route[0]
    x = -x 
    y = -y 
    tmp_route = []
    # 돌진할 산타와 거리가 가까워지는 방향 찾기
    for i in range(8):
        nx = r_x + dx[i]
        ny = r_y + dy[i]
        new_dist = (nx-x)**2 + (ny-y)**2 
        if nx < 0 or nx >=n or ny < 0 or ny >= n :
            continue 
            
        if new_dist < dist :
            tmp_route.append((new_dist, i, nx, ny))
    
    tmp_route.sort()
    # 거리, 이동 방향, 이동 후 위치
    dist, direct, x, y = tmp_route[0]
    # 루돌프 위치 갱신 
    r_x = x 
    r_y = y 
    
    # 루돌프가 이동한 위치에 산타가 있으면 밀어낸다 
    if graph[r_x][r_y]:
        p_id = graph[r_x][r_y]
        # 밀어낸 위치 
        nx = r_x + dx[direct] * c 
        ny = r_y + dy[direct] * c 
        # 밀어내기 
        graph[r_x][r_y] = 0
        # 기절 시간 갱신 
        stop_time[p_id] = t + 1 
        # 점수 더하기
        scores[p_id] += c 
        shoot(p_id, nx, ny, direct)

2. 산타 이동 함수

s_dx = [-1,0,1,0]
s_dy = [0,1,0,-1]

def move_santa():
    global r_x, r_y 

    for i in range(1,p+1):
        # 범위에 없는 산타 
        if out[i] : 
            continue 
        # 기절한 산타 
        if stop_time[i] >= t :
            continue 
        x, y = santa[i]
        route = []

        next_x = x 
        next_y = y 
        p_d = -1 
        # 루돌프와 가까워지는 방향 구하기 
        for j in range(4):
            nx = x + s_dx[j]
            ny = y + s_dy[j]
            origin = (x-r_x)**2 + (y-r_y)**2 
            if nx <0 or nx >= n or ny < 0 or ny >= n :
                continue 
            # 이미 산타가 있으면 이동할 수 없음 
            if graph[nx][ny] :
                continue 
            
            new_dist = (nx-r_x)**2 + (ny-r_y)**2

            if new_dist < origin :
                route.append((new_dist, j, nx, ny))
        
        # 가까워지는 방향으로 갈 수 없음 -> 움직이지 않음 
        if len(route) == 0 :
            continue 

        route.sort()
        # 거리, 이동한 방향, 이동한 위치 
        dist, p_d, next_x, next_y = route[0]
        # 이동한 위치에 루돌프가 있으면 
        if r_x == next_x and r_y == next_y :
            # 이동한 방향과 반대방향으로 밀어남 
            p_d = (p_d+2)%4 
            nx = r_x + s_dx[p_d] * d
            ny = r_y + s_dy[p_d] * d 
            # 밀어내기 
            graph[x][y] = 0 
            # 점수 더하기 
            scores[i] += d 
            # 기절 시간 갱신 
            stop_time[i] = t + 1 
            shoot(i, nx, ny, p_d)
        # 빈칸이면 바로 이동 
        else :
            graph[x][y] = 0 
            graph[next_x][next_y] = i
            santa[i] = [next_x,next_y]

3. 매 턴 이후 점수 더하는 함수 

def get_point():
    cnt = 0 

    for i in range(1,p+1):
        # 범위에 없는 산타
        if out[i]:
            continue 
        # 점수 더하기 
        scores[i] += 1 
        cnt += 1 
    
    return cnt

 

전체 코드는 아래와 같다.

n, m, p, c, d = map(int,input().split())
r_x, r_y = map(int,input().split())
r_x -=1 
r_y -=1

# index에 해당하는 산타의 위치 저장 
santa = [ [] for _ in range(p+1) ] 
# 위치별 산타 번호 저장 
graph = [ [0] * n for _ in range(n) ]
# index에 해당하는 산타의 점수 저장 
scores = [0] * (p+1)
# index에 해당하는 산타 존재 여부 저장 
out = [False] * (p+1)
# index에 해당하는 산타가 언제까지 기절하는지 시간 저장 
stop_time = [0] * (p+1)
# 몇 초가 흘렸는지 
t = 1 

for _ in range(p):
    i, a, b = map(int,input().split())
    a -= 1 
    b -= 1 
    santa[i] = [a,b]
    graph[a][b] = i 

# 상하좌우대각선 8방향 
dx = [-1,0,1,0,1,1,-1,-1]
dy = [0,1,0,-1,1,-1,1,-1]

# 산타 번호, 산타 위치, 방향 
def shoot(s_id, x, y, direct):
    # 밀어내진 위치가 범위를 벗어나면 out 
    if x < 0 or x >= n or y < 0 or y >= n:
        out[s_id] = True 
        return 
    
    # 밀어난 위치에도 산타가 있으면 재귀적으로 밀어낸다 
    if graph[x][y] :
        tmp_id = graph[x][y]
        # 착지 
        graph[x][y] = s_id
        santa[s_id] = [x,y]
        # 연쇄적으로 밀어냄 
        shoot(tmp_id, x+dx[direct], y+dy[direct], direct)
    # 밀어난 위치가 빈 칸이라면 바로 착지 
    else :
        graph[x][y] = s_id
        santa[s_id] = [x,y]

def move_rudolf():
    global r_x, r_y 

    # 산타 별 거리 구하기
    route = []
    for i in range(1,p+1):
        if out[i]:
            continue 

        s_x, s_y = santa[i]
        dist = (s_x-r_x)**2 + (s_y-r_y)**2 
        # 거리 작은 순 -> r 큰 순 -> c 큰 순으로 정렬 
        route.append([dist, -s_x, -s_y, i])

    route.sort()
    
    # 거리, 돌진할 산타 위치, 돌진한 산타 번호
    dist, x, y, i = route[0]
    x = -x 
    y = -y 
    tmp_route = []
    # 돌진할 산타와 거리가 가까워지는 방향 찾기
    for i in range(8):
        nx = r_x + dx[i]
        ny = r_y + dy[i]
        new_dist = (nx-x)**2 + (ny-y)**2 
        if nx < 0 or nx >=n or ny < 0 or ny >= n :
            continue 
            
        if new_dist < dist :
            tmp_route.append((new_dist, i, nx, ny))
    
    tmp_route.sort()
    # 거리, 이동 방향, 이동 후 위치
    dist, direct, x, y = tmp_route[0]
    # 루돌프 위치 갱신 
    r_x = x 
    r_y = y 
    
    # 루돌프가 이동한 위치에 산타가 있으면 밀어낸다 
    if graph[r_x][r_y]:
        p_id = graph[r_x][r_y]
        # 밀어낸 위치 
        nx = r_x + dx[direct] * c 
        ny = r_y + dy[direct] * c 
        # 밀어내기 
        graph[r_x][r_y] = 0
        # 기절 시간 갱신 
        stop_time[p_id] = t + 1 
        # 점수 더하기
        scores[p_id] += c 
        shoot(p_id, nx, ny, direct)

s_dx = [-1,0,1,0]
s_dy = [0,1,0,-1]

def move_santa():
    global r_x, r_y 

    for i in range(1,p+1):
        # 범위에 없는 산타 
        if out[i] : 
            continue 
        # 기절한 산타 
        if stop_time[i] >= t :
            continue 
        x, y = santa[i]
        route = []

        next_x = x 
        next_y = y 
        p_d = -1 
        # 루돌프와 가까워지는 방향 구하기 
        for j in range(4):
            nx = x + s_dx[j]
            ny = y + s_dy[j]
            origin = (x-r_x)**2 + (y-r_y)**2 
            if nx <0 or nx >= n or ny < 0 or ny >= n :
                continue 
            # 이미 산타가 있으면 이동할 수 없음 
            if graph[nx][ny] :
                continue 
            
            new_dist = (nx-r_x)**2 + (ny-r_y)**2

            if new_dist < origin :
                route.append((new_dist, j, nx, ny))
        
        # 가까워지는 방향으로 갈 수 없음 -> 움직이지 않음 
        if len(route) == 0 :
            continue 

        route.sort()
        # 거리, 이동한 방향, 이동한 위치 
        dist, p_d, next_x, next_y = route[0]
        # 이동한 위치에 루돌프가 있으면 
        if r_x == next_x and r_y == next_y :
            # 이동한 방향과 반대방향으로 밀어남 
            p_d = (p_d+2)%4 
            nx = r_x + s_dx[p_d] * d
            ny = r_y + s_dy[p_d] * d 
            # 밀어내기 
            graph[x][y] = 0 
            # 점수 더하기 
            scores[i] += d 
            # 기절 시간 갱신 
            stop_time[i] = t + 1 
            shoot(i, nx, ny, p_d)
        # 빈칸이면 바로 이동 
        else :
            graph[x][y] = 0 
            graph[next_x][next_y] = i
            santa[i] = [next_x,next_y]

def get_point():
    cnt = 0 

    for i in range(1,p+1):
        # 범위에 없는 산타
        if out[i]:
            continue 
        # 점수 더하기 
        scores[i] += 1 
        cnt += 1 
    
    return cnt 

# for debug
def print_graph():
    global r_x, r_y
    for i in graph:
        print(i)
    print(santa)
    print(out)
    print(scores)
    print(r_x, r_y)

#print_graph()
while t <= m :
    move_rudolf()
    #print_graph()
    move_santa()
    #print_graph()
    cnt = get_point()
    #print_graph()
    # 범위에 존재하는 산타가 없으면 바로 종료
    if cnt == 0 :
        break 
    t += 1 
    #print('---------------')

print(*scores[1:])