문제

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

나의 코드

"""

"""

N = int(input())
pan = [ [0] * N for _ in range(N) ]
stu = [ list(map(int, input().split())) for _ in range(N**2) ]

dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]
def seat_arrange(st): # 학생을 배치하는 함수
    st_num, like = st[0], st[1:] # 학생의 번호와, 좋아하는 학생의 번호 목록
    choice_list = [] # 인접한 칸에 (좋아하는 학생의 수, 빈 칸의 수, 현재 좌표(x, y)) 를 담는 리스트

    for i in range(N):
        for j in range(N):
            like_space = 0 # 좋아하는 학생의 수
            empty_space = 0 # 빈 칸의 수
            for k in range(4):
                mx = i + dx[k]
                my = j + dy[k]

                if not (0 <= mx < N and 0 <= my < N):
                    continue

                if pan[mx][my] in like: # 인접한 칸에 좋아하는 학생이 있다면 like_space에 +1
                    like_space += 1
                elif pan[mx][my] == 0: # 인접한 칸이 비어있는 칸이라면 empty_space에 +1
                    empty_space += 1

            choice_list.append((like_space, empty_space, i, j))

    # 좋아하는 학생(내림차순) -> 빈 칸의 수(내림차순) -> 행의 번호(오름차순) -> 열의 번호(오름차순) 으로 정렬
    choice_list = sorted(choice_list, key = lambda x : (x[0], x[1], -x[2], -x[3]), reverse=True)

    for _, _, x, y in choice_list: # 1순위 좌표가 비어있다면 학생을 배치하고 이미 다른 학생이 있다면 다음 우선순위의 위치에 배치
        if pan[x][y] == 0:
            pan[x][y] = st_num
            break

for st in stu:
    seat_arrange(st)

def satisfaction(st): # 인접한 곳에 좋아하는 친구 수에 따라 만족도를 구하는 함수
    global answer

    st_num, like = st[0], st[1:] # 학생의 번호와 좋아하는 학생의 번호 목록

    for i in range(N):
        for j in range(N):
            if pan[i][j] == st_num:
                cnt = 0
                for k in range(4): # 좋아하는 학생의 수를 세어 만족도를 센다
                    mx = i + dx[k]
                    my = j + dy[k]

                    if not (0 <= mx < N and 0 <= my < N):
                        continue

                    if pan[mx][my] in like:
                        cnt += 1

                answer += satis[cnt]

satis = [0, 1, 10, 100, 1000] # 만족도 리스트
answer = 0
for st in stu:
    satisfaction(st)

print(answer)

    

설명

람다식을 이용한 내림차순+오름차순 정렬이 핵심 (단순 구현 문제)

참고 자료