[코드트리] 2025 하반기 오후 1번 - AI 로봇청소기

정말 오랜만에 푼 알고리즘 문제인데, 2d list로 해결해야 할 문제를 dict로 접근해서 시간초과 문제가 있었다. 

dict, list 시간복잡도를 잘 알아둬야겠다. 

코딩테스트 기출 문제 제출: AI 로봇청소기 | 코드트리

from collections import deque

# 1. 입력 값
N, K, L = map(int, input().split())
board = [] # 장애물 탐색용
cleaner = [] # 로청 위치

for i in range(N):
    line = list(map(int, input().split()))
    board.append(line)

for i in range(K):
    r, c = list(map(int, input().split()))
    cleaner.append([r-1, c-1])

def print_board(clner, bd):
    for i in range(N):
        for j in range(N):
            if [i, j] in clner:
                print(f"C/{bd[i][j]:>2}", end=" ")
            else:
                target = bd[i][j]
                print(f"{target:>4}", end=" ")
        print()
    print()

def bfs(clner, x, y, bd):
    # 가장 가까운 오염격자로 이동
    clner = set(tuple(c) for c in clner)
    visited = [[False] * N for _ in range(N)]

    visited[x][y] = True
    queue = deque([(x, y, 0)])
    candidate = []
    min_dist = float("inf")

    while queue:
        tmpx, tmpy, dist = queue.popleft()

        if dist > min_dist:
            break

        if bd[tmpx][tmpy] > 0 and (tmpx, tmpy) not in clner:
            if min_dist == float("inf"):
                min_dist = dist
            if dist == min_dist:
                candidate.append([tmpx, tmpy])
            continue
        
        for dx, dy in [[tmpx, tmpy+1], [tmpx, tmpy-1], [tmpx+1, tmpy], [tmpx-1, tmpy]]:
            if dx < 0 or dx >= N or dy < 0 or dy >= N or visited[dx][dy] or bd[dx][dy] == -1 or (dx, dy) in clner:
                continue
            visited[dx][dy] = True
            queue.append((dx, dy, dist+1))

    return candidate
    

def move(clner, bd):
    # 이동 거리가 가장 가까운 오염된 격자로 이동
    for i in range(len(clner)):
        r, c = clner[i]

        if bd[r][c] == 0:
            candidate = bfs(clner, r, c, bd)
            candidate.sort(key=lambda x: (x[0], x[1]))
            if candidate:
                clner[i] = candidate[0]

    return clner


def clean(clner, bd):
    for x, y in clner:
        delta = [[0, 0], [-1, 0], [1, 0], [0, -1], [0, 1]] # 나, 위, 아래, 왼, 오
        direction = [3, 1, 4, 2]
        max_direction = -1
        max_value = -1

        # 1. 먼지량 가장 큰 방향 결정
        for d in direction:
            tmp_value = 0
            for i in range(len(delta)):
                next_x, next_y = x+delta[i][0], y+delta[i][1]
                if d != i and next_x >= 0 and next_x < N and next_y >= 0 and next_y < N and bd[next_x][next_y] > 0:
                    tmp_value += min(bd[next_x][next_y], 20)
            if tmp_value > max_value:
                max_direction = d
                max_value = tmp_value

        # 2. 청소
        for i in range(len(delta)):
            next_x, next_y = x+delta[i][0], y+delta[i][1]
            if max_direction != i and next_x >= 0 and next_x < N and next_y >= 0 and next_y < N \
                and bd[next_x][next_y]>0:
                bd[next_x][next_y] = max(0, bd[next_x][next_y]-20)
    return bd

def collect(clner, bd):
    # 먼지가 있는 모든 격자에 동시에 5씩 추가
    for i in range(N):
        for j in range(N):
            if bd[i][j] > 0:
                bd[i][j] += 5
    return bd

def spread(clner, bd):
    # 깨끗한 격자에 주변 4방향 격자의 먼지량 합을 10으로 나눈 값만큼 먼지가 확산됨
    import copy
    next_bd = copy.deepcopy(bd)
    for i in range(N):
        for j in range(N):
            if bd[i][j] == 0:
                for x, y in [[i, j+1], [i, j-1], [i+1, j], [i-1, j]]:
                    if x < 0 or x >= N or y < 0 or y >= N or bd[x][y] == -1: 
                        continue
                    next_bd[i][j] += bd[x][y]
                next_bd[i][j] //= 10

    return next_bd

def answer(bd):
    value = 0
    for i in range(N):
        for j in range(N):
            if bd[i][j] > 0:
                value += bd[i][j]
    return value

# print_board(cleaner, board)
for tst in range(L):

    # 1. 청소기 이동
    cleaner = move(cleaner, board)
    # print("이동 완료")
    # print_board(cleaner, board)


    # 2. 청소
    board = clean(cleaner, board)
    # print("청소 완료")
    # print_board(cleaner, board)
    

    # 3. 먼지 축적
    board = collect(cleaner, board)
    # print("먼지 축적 완료")
    # print_board(cleaner, board)


    # 4. 먼지 확산
    board = spread(cleaner, board)
    # print("확산 완료")
    # print_board(cleaner, board)

    # 5. 출력
    print(answer(board))