[백준] 삼성 SW 역량 테스트 기출 문제21608번 : 상어 초등학교 (파이썬/Python)

[백준] 삼성 SW 역량 테스트 기출 문제 21608번 : 상어 초등학교 (파이썬/Python)

정답 코드 및 문제 설명

정답 전체 코드

n = int(input()) # 교실 크기 
data = [[0] * n for _ in range(n)] # 교실 격자 초기화 
students = [list(map(int, input().split())) for _ in range(n**2)] # 각 학생과 좋아하는 학생 4명 번호 저장

#상하 좌우 움직일때 좌표 미리 저장
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: #현재 선택 격자가 비었다면
                # 각 칸마다 좋아하는 학생수, 빈자리 수를 계산
                prefer, empty = 0, 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 student[1:]:
                            prefer += 1
                        if data[nx][ny] == 0:
                            empty += 1
                #계산된 결과를 available 리스트에 저장된다.
                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
#만족도 점수 배열 0~4명 좋아하는 학생 수에 따라 다음
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][1:]:
                    count += 1
        answer += score[count] 
print(answer)

문제 및 코드 설명

문제에는 아래 3가지 조건을 순서대로 우선순위를 가진다.
1. 비어있는 칸 중에서 좋아하는 학생이 인접한 칸에 가장 많은 칸으로 자리를 정한다.
2. 1을 만족하는 칸이 여러 개이면, 인접한 칸 중에서 비어있는 칸이 가장 많은 칸으로 자리를 정한다.
3. 2를 만족하는 칸도 여러 개인 경우에는 행의 번호가 가장 작은 칸으로, 그러한 칸도 여러 개이면 열의 번호가 가장 작은 칸으로 자리를 정한다.

즉, 좋아하는 학생이 인접한 칸에 얼마나 있는지, 인접한 칸 중 빈칸, 행, 열 순으로 정렬해야한다.

available.sort(key=lambda x: (-x[2], -x[3], x[0], x[1]))

코드 내에 sort 가 쓰인 부분을 보면 알 수 있다.
available 리스트는 각 학생마다 [행 , 열, 좋아하는 학생 수, 인접한 빈칸 수] 로 저장되어있다.

람다 함수를 사용하여 우선 순위를 설정하여 sort, 정렬할 수 있다.
x: -x[2]
와 같이 - 가 붙어있으면 내림 차순으로, x[0] 처럼 붙어있지 않다면 해당하는 것을 오름차순으로 정렬한다.

즉, 위 코드는 좋아하는 학생 수 x[2] 를 내림차순으로 먼저 정렬하고, 인접한 빈칸 수 x[3] 을 내림차순으로 그 다음으로 행 과 열 을 오름차순으로 정렬한다.