코드트리 싸움땅 (파이썬)

https://www.codetree.ai/training-field/frequent-problems/problems/battle-ground

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

import sys

input = sys.stdin.readline

# 상우하좌 순서
dx = [-1, 0, 1, 0]
dy = [0, 1, 0, -1]

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

# 총이 없는 칸은 [], 총이 있는 칸은 [n] 으로 a에 저장
a = []
for i in range(n):
    a.append([])
    temp = list(map(int, input().split()))
    for j in range(n):
        if temp[j] == 0:
            a[i].append([])
        else:
            a[i].append([temp[j]])

# player: 각 플레이어의 좌표, 이동방향, 초기점수, 소지한 총의 공격력
# number: 해당 좌표에 몇 번째 플에이어가 있는지
player = []
number = [[-1 for _ in range(n)] for _ in range(n)]
for i in range(m):
    x, y, d, s = map(int, input().split())
    player.append([x-1, y-1, d, s, 0])
    number[x-1][y-1] = i


# 두 플레이어가 싸우는 경우
# low_idx, high_idx: 패배/승리한 플레이어 인덱스
# nx, ny: 두 플레이어가 마주친 좌표
def fight(low_idx, high_idx, nx, ny):
    global ans
    lx, ly, ld, ls, lg = player[low_idx]
    hx, hy, hd, hs, hg = player[high_idx]
    # 승리한 플레이어 점수 추가
    ans[high_idx] += (hs + hg) - (ls + lg)
    # 패배한 플레이어 총 버리기
    if lg > 0:
        a[nx][ny].append(lg)
    lg = 0
    # 시계방향으로 방향을 바꾸면서 이동가능한 칸 찾기
    while True:
        nlx = lx + dx[ld]
        nly = ly + dy[ld]
        if not 0 <= nlx < n or not 0 <= nly < n:
            ld = (ld + 1) % 4
            continue
        if number[nlx][nly] != -1:
            ld = (ld + 1) % 4
            continue
        # 해당 칸에 총이 있으면, 가장 공격력이 높은 총 획득
        if a[nlx][nly]:
            a[nlx][nly] = sorted(a[nlx][nly], reverse=True)
            lg, a[nlx][nly] = a[nlx][nly][0], a[nlx][nly][1:]
        # 패배한 플레이어 player, number 수정
        player[low_idx] = [nlx, nly, ld, ls, lg]
        number[nlx][nly] = low_idx
        break
    # 승리한 플레이어는 가장 공격력이 큰 총으로 교체
    a[nx][ny] = sorted(a[nx][ny], reverse=True)
    if a[nx][ny]:
        if a[nx][ny][0] > hg:
            if hg > 0:
                a[nx][ny].append(hg)
            hg, a[nx][ny] = a[nx][ny][0], a[nx][ny][1:]
    # 승리한 플레이어 player, number 수정
    player[high_idx] = [nx, ny, hd, hs, hg]
    number[nx][ny] = high_idx


ans = [0 for _ in range(m)]
for _ in range(k):
    for i, p in enumerate(player):
        x, y, d, s, g = p
        nx = x + dx[d]
        ny = y + dy[d]
        # 범위를 벗어나면 반대 방향으로 수정
        if not 0 <= nx < n or not 0 <= ny < n:
            d = (d + 2) % 4
            nx = x + dx[d]
            ny = y + dy[d]
            player[i][2] = d
        # 이동 전 좌표에 대한 number 초기화
        number[x][y] = -1
        # 이동한 칸에 플레이어가 없는 경우
        if number[nx][ny] == -1:
            # 총이 있으면, 가장 공격력이 높은 총으로 교체
            if a[nx][ny]:
                a[nx][ny] = sorted(a[nx][ny], reverse=True)
                if a[nx][ny][0] > g:
                    if g > 0:
                        a[nx][ny].append(g)
                    g, a[nx][ny] = a[nx][ny][0], a[nx][ny][1:]
            # 이동 후 바뀐 값으로 player, number 수정
            player[i] = [nx, ny, d, s, g]
            number[nx][ny] = i
        # 이동한 칸에 플레이어가 있는 경우
        else:
            # 해당 칸에 먼저 있던 플레이어의 인덱스와 초기값, 총 공격력 가져오기
            bi = number[nx][ny]
            _, _, _, bs, bg = player[bi]
            # 이동한 후 좌표만 수정
            player[i][0], player[i][1] = nx, ny
            # 문제 조건대로 패배/승리 플레이어의 인덱스를 fight 함수에 입력
            if s + g > bs + bg:
                fight(bi, i, nx, ny)
            elif s + g < bs + bg:
                fight(i, bi, nx, ny)
            else:
                if s > bs:
                    fight(bi, i, nx, ny)
                else:
                    fight(i, bi, nx, ny)

for n in ans:
    print(n, end=' ')