[백준: Python] 21608 - 상어 초등학교

[문제]

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

기본적인 구현문제로, 주어진 조건대로 착실하게 수행하면 되는 문제입니다.

---

[초기화]

    def __init__(self):
        self.n = int(input())
        self.student_info = [tuple(map(int, input().split())) for _ in range(self.n**2)]
        self.score = {0: 0, 1: 1, 2: 10, 3: 100, 4: 1000}
        self.grid = [[0] * self.n for _ in range(self.n)]
        self.friends_list = defaultdict(set)
        self.answer = 0

self.score의 경우, 인접한 친구수에 따른 점수표입니다. 직관적으로 확인하기 위해 딕셔너리 구조를 사용했지만, 순서가 0번부터 순차적이므로 리스트로 관리해도 무방합니다.

 

self.friends_list = defaultdict(set)의 경우, 주어진 학생에 따른 좋아하는 학생을 딕셔너리 구조로 관리하기 위한 데이터입니다. 값에 해당하는 자료구조를 집합으로 사용한 이유는 인접한 학생이 좋아하는 학생인지 확인할 때 빠르게 처리하기 위함입니다.

---

[풀이]

- 입력으로 주어지는 학생과 해당 학생이 선호하는 학생들 분리

    def init_structure(self):
        for i in self.student_info:
            student, like = i[0], set(i[1:])
            self.friends_list[student] = like

입력으로 주어지는 학생 순서대로 키로 두고, 좋아하는 학생들을 집합으로 변환해 값으로 저장합니다.

 

- 학생의 자리 탐색

    def search(self):
        for i in self.student_info:
            student = i[0]
            max_friend_cnt = max_empty_cnt = -1
            for y in range(self.n):
                for x in range(self.n):
                    if self.grid[y][x] != 0:
                        continue
                    friend_cnt = empty_cnt = 0
                    for dy, dx in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                        my, mx = y + dy, x + dx
                        if 0 <= my < self.n and 0 <= mx < self.n:
                            if self.grid[my][mx] == 0:
                                empty_cnt += 1
                            if self.grid[my][mx] in self.friends_list[student]:
                                friend_cnt += 1

                    if max_friend_cnt < friend_cnt or (max_friend_cnt == friend_cnt and max_empty_cnt < empty_cnt):
                        max_friend_cnt = friend_cnt
                        max_empty_cnt = empty_cnt
                        ey, ex = y, x

            self.grid[ey][ex] = student

입력으로 주어진 학생들의 자리를 탐색합니다.

 

3번의 조건을 만족하기 위해 첫 인덱스부터 순차적으로 탐색을 수행하면 자리를 확인합니다.

1. 빈자리가 아니라면(값이 0이 아닌 경우), 다음 자리 탐색
2. 상하좌우로 인접한 자리를 확인하며, 해당 자리가 빈자리인지, 좋아하는 학생이 앉아있는 확인 하며 각각 카운트합니다.
3. 인접한 자리에 대한 탐색이 끝났다면, 최적의 자리를 찾기 위해 카운트한 값을 갱신해 나가며 적절한 위치를 저장합니다.
   • 지금까지 확인한 자리 중 좋아하는 학생이 가장 많거나, 좋아하는 학생이 동일하나 빈자리가 많을 경우 값 갱신
   • 최적의 위치 갱신
   • 빈자리가 많은 곳은 선택하는 이유는 빈자리가 많은수록 인접한 자리에 좋아하는 학생을 많이 앉힐 수 있기 때문
4. 최적의 자리를 찾았다면, 해당 위치에 입력으로 주어진 학생을 앉힘

 

- 점수 갱신

    def get_score(self):
        for y in range(self.n):
            for x in range(self.n):
                cnt = 0
                for dy, dx in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                    my, mx = y + dy, x + dx
                    student = self.grid[y][x]
                    if 0 <= my < self.n and 0 <= mx < self.n and self.grid[my][mx] in self.friends_list[student]:
                        cnt += 1

                self.answer += self.score[cnt]

모든 학생을 다 앉혔다면, 순차적으로 탐색하며 앉은 학생 주변에 좋아하는 학생이 얼마나 앉아있는지 확인해 좋아하는 학생수에 따라 정답을 갱신합니다.

 

- 코드 실행 순서

    def solve(self):
        self.init_structure()
        self.search()
        self.get_score()

        print(self.answer)

---

[전체 코드]

from collections import defaultdict


class Main:
    def __init__(self):
        self.n = int(input())
        self.student_info = [tuple(map(int, input().split())) for _ in range(self.n**2)]
        self.score = {0: 0, 1: 1, 2: 10, 3: 100, 4: 1000}
        self.grid = [[0] * self.n for _ in range(self.n)]
        self.friends_list = defaultdict(set)
        self.answer = 0

    def init_structure(self):
        for i in self.student_info:
            student, like = i[0], set(i[1:])
            self.friends_list[student] = like

    def search(self):
        for i in self.student_info:
            student = i[0]
            max_friend_cnt = max_empty_cnt = -1
            for y in range(self.n):
                for x in range(self.n):
                    if self.grid[y][x] != 0:
                        continue
                    friend_cnt = empty_cnt = 0
                    for dy, dx in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                        my, mx = y + dy, x + dx
                        if 0 <= my < self.n and 0 <= mx < self.n:
                            if self.grid[my][mx] == 0:
                                empty_cnt += 1
                            if self.grid[my][mx] in self.friends_list[student]:
                                friend_cnt += 1

                    if max_friend_cnt < friend_cnt or (max_friend_cnt == friend_cnt and max_empty_cnt < empty_cnt):
                        max_friend_cnt = friend_cnt
                        max_empty_cnt = empty_cnt
                        ey, ex = y, x

            self.grid[ey][ex] = student

    def get_score(self):
        for y in range(self.n):
            for x in range(self.n):
                cnt = 0
                for dy, dx in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                    my, mx = y + dy, x + dx
                    student = self.grid[y][x]
                    if 0 <= my < self.n and 0 <= mx < self.n and self.grid[my][mx] in self.friends_list[student]:
                        cnt += 1

                self.answer += self.score[cnt]

    def solve(self):
        self.init_structure()
        self.search()
        self.get_score()

        print(self.answer)


problem = Main()
problem.solve()

---