[백준] 21608_상어 초등학교 (Python)

📕 문제

https://www.acmicpc.net/problem/21608

21608번: 상어 초등학교

 

💡 풀이법

1. 학생을 키로 하고 학생이 선호하는 학생의 리스트를 value로 하는 like 딕셔너리를 만들었다.

2. 학생들을 위치시킬 N*N의 2차 행렬 room을 만들었다.

3. 각 학생마다 교실을 돌면서 각 자리의 행과 열, 주변의 좋아하는 사람 수, 빈 자리 수 를 구해서 priority리스트에 넣는다.

4. 모든 자리를 다 돌고 난 후 priority를 정렬해서 가장 우선순위가 높은 자리를 구해서 그 자리에 학생을 배치한다.

5. 모든 학생을 배치한 후 다시 한번 교실을 돌면서 만족도를 계산한다.

 

N = int(input())
like = {}
for _ in range(N**2):
    lst = list(map(int, input().split()))
    like[lst[0]] = lst[1:]

room = [[0 for _ in range(N)] for _ in range(N)]
dxy = [(1,0),(-1,0),(0,1),(0,-1)]

for key in like:
    priority = []
    for x in range(N):
        for y in range(N):
            like_cnt, empty_cnt = 0, 0
            if not room[x][y]:
                for dx, dy in dxy:
                    nx, ny = x + dx, y + dy
                    if 0 <= nx < N and 0 <= ny < N:
                        if not room[nx][ny]:
                            empty_cnt += 1
                        elif room[nx][ny] in like[key]:
                            like_cnt += 1
                priority.append((like_cnt, empty_cnt, x, y))
    priority.sort(key=lambda x:(-x[0],-x[1],x[2],x[3]))
    room[priority[0][2]][priority[0][3]] = key

answer = 0
for x in range(N):
    for y in range(N):
        like_cnt = 0
        for dx, dy in dxy:
            nx, ny = x + dx, y + dy
            if 0 <= nx < N and 0 <= ny < N:
                if room[nx][ny] in like[room[x][y]]:
                    like_cnt += 1
        if like_cnt:
            answer += 10**(like_cnt-1)
print(answer)