[백준/파이썬] #21608

30분이 초과했다.... 함수 단위로 분해하려고 했는데 생각보다 복잡해져서 꼬여버린 것 같다. 

import heapq
from collections import defaultdict

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


def best_seat(liked_students):
    global adjacent_vacant_seat
    adjacent_seats = defaultdict(lambda: 0)
    for liked_student in liked_students:

        # 좋아하는 학생이 아직 자리를 배정받지 않았을 경우
        if not student_to_seat[liked_student]:
            continue

        current_row, current_col = student_to_seat[liked_student]
        for k in range(4):
            nrow, ncol = current_row + dx[k], current_col + dy[k]
            if nrow < 0 or nrow >= N or ncol < 0 or ncol >= N:
                continue

            # 인접한 좋아하는 학생 수 구하기
            adjacent_seats[(nrow, ncol)] += 1

    if len(adjacent_seats) == 0:
        adjacent_count, row, col = heapq.heappop(adjacent_vacant_seat)
        if adjacent_count > 0:
            heapq.heappush(adjacent_vacant_seat, (adjacent_count-1, row, col))


def sit():
    global seat, student_to_seat
    for s in students:
        student, liked_students = s[0], s[1:]
        row, col = best_seat(liked_students)
        seat[row][col] = student
        student_to_seat[student] = (row, col)


def calculate_satisfaction():
    global answer
    pass


if __name__ == "__main__":
    N = int()
    answer = 0
    students = []
    student_to_seat = defaultdict(list)
    seat = [[0] * N for _ in range(N)]
    adjacent_vacant_seat = []

    for i in range(N*N):
        s = list(map(int, input().split()))
        students.append(s)
        student_to_seat[s[0]] = False

    for i in range(N):
        for j in range(N):
            seat[i][j] = 0
            if (i == 0 or i == N-1) and (j == 0 or j == N-1):
                heapq.heappush(adjacent_vacant_seat, (2, i, j))
            elif (i == 0 or i == N-1) or (j == 0 or j == N-1):
                heapq.heappush(adjacent_vacant_seat, (3, i, j))
            else:
                heapq.heappush(adjacent_vacant_seat, (4, i, j))

    sit(students)
    calculate_satisfaction()
    print(answer)

 

다른 풀이를 참고해서 풀었다. 

N = int(input())
students = [list(map(int, input().split())) for _ in range(N * N)]
data = [[0] * N for _ in range(N)]

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

for student in students:
    available = []

    for i in range(N):
        for j in range(N):
            if data[i][j] != 0:
                continue
            
            prefer, empty = 0, 0
            for k in range(4):
                nx, ny = i + dx[k], j + dy[k]
                if nx < 0 or nx >= N or ny < 0 or ny >= N:
                    continue

                if data[nx][ny] in student[1:]:
                    prefer += 1

                if data[nx][ny] == 0:
                    empty += 1
            
            available.append((i, j, prefer, empty))
    
    available.sort(key=lambda x: (-x[2], -x[3], x[0], x[1]))
    data[available[0][0]][available[0][1]] = student[0]


answer = 0
score = [0, 1, 10, 100, 1000]
students.sort()

for i in range(N):
    for j in range(N):
        count = 0

        for k in range(4):
            nx, ny = i + dx[k], j + dy[k]
            if 0 <= nx < N and 0 <= ny < N:
                if data[nx][ny] in students[data[i][j] - 1]:
                    count += 1

        answer += score[count]

print(answer)