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

문제

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

21608번: 상어 초등학교

문제에서 자리 배정하는 과정을 그림으로 보여주고 있으니 보고오자

 

접근법

N이 최대 20이기 때문에 이중 for문을 써도 괜찮다. 따라서 학생마다 모든 자리를 돌아보며 비어있는 자리를 우선순위 큐에 넣었고, 마지막에 하나만 꺼내서 자리를 배정했다. 자리를 나타내는 class를 만들고, compareTo() 함수를 오버라이딩해서 문제에서 제시하는 우선순위를 그대로 나타내주었다.  

 

시간을 조금이라도 줄이기 위해 학생이 좋아하는 학생들의 번호를 HashMap을 사용해 저장하였다. HashMap은 순서가 보장되지 않으므로 자리를 배치하는 순서는 배열을 사용해 따로 저장해주었다.

 

처음에는 자리를 모두 배치한 뒤, 만족도를 한번에 계산하면 되지 않을까 했는데 시간이 조금이라도 줄어들까 싶어 자리를 배치하는 과정에서 만족도까지 구해주었다. 실제로 시간이 줄어드는지는 모르겠다..!

 

알고리즘

코드를 작성하기 전에 글로 적은 알고리즘으로 이해하고 코드를 작성하면 도움이 될 것 같다.

<input>
1. 학생의 번호를 int 배열에 저장한다.
2. 학생의 번호를 key, 해당 학생이 좋아하는 학생들을 value로 HashMap에 저장한다.

<자리 배정 및 만족도>
1. 순서를 저장해놓은 int 배열에서 학생의 번호를 가져온다.
2. 모든 자리를 돌아보면서 아직 배정되지 않은 자리가 있으면 해당 자리의 정보(비어있는 자리의 수, 좋아하는 학생의 수)를 가져와 우선순위 큐에 넣는다.
3. 큐에서 한 자리를 꺼내서 배정한다.
4. 해당 자리에서 인접한 자리에 앉아있는 좋아하는 학생의 수는 이미 위에서 계산했으므로 그대로 저장해준다.
5. 나와 인접한 자리에 앉은 학생 중, 나를 좋아하는 학생이 있다면 그 학생의 만족도를 +1 해준다.
6. 모든 학생의 자리 배정이 끝나면 최종적으로 만족도를 계산한다.

 

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.*;

public class Main {

    static int[] dirR = {-1, 1, 0, 0};
    static int[] dirC = {0, 0, -1, 1};

    static int N;
    static int[] order;                             // 학생의 순서
    static HashMap<Integer, Set<Integer>> students; // 학생의 정보
    static int[][] classroom;                       // 교실의 자리 배치도
    static int[][] satisfaction;                    // 각 자리별 만족도

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        N = Integer.parseInt(br.readLine());

        order = new int[N * N];
        students = new HashMap<>();
        classroom = new int[N][N];
        satisfaction = new int[N][N];

        for (int i = 0; i < N * N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int n = Integer.parseInt(st.nextToken()); // 학생의 번호
            order[i] = n;

            Set<Integer> likeFriends = new HashSet<>();
            for (int j = 0; j < 4; j++) {
                likeFriends.add(Integer.parseInt(st.nextToken()));
            }

            students.put(n, likeFriends);
        } // 입력 끝

        // 1. 자리 배정
        for (int i = 0; i < N * N; i++) {
            PriorityQueue<Seat> pq = new PriorityQueue<>();
            int number = order[i];

            // 비어있는 자리가 있으면 해당 자리의 정보를 우선 순위 큐에 넣는다.
            for (int r = 0; r < N; r++) {
                for (int c = 0; c < N; c++) {
                    if (classroom[r][c] == 0) {
                        pq.offer(seatInfo(r, c, students.get(number)));
                    }
                }
            }

            Seat seat = pq.poll(); // 배정된 자리
            classroom[seat.r][seat.c] = number;
            satisfaction[seat.r][seat.c] = seat.like;

            // 나와 인접한 자리에 앉은 학생 중, 나를 좋아하는 학생이 있으면 만족도를 1 올려준다.
            for (int j = 0; j < 4; j++) {
                int r = seat.r + dirR[j];
                int c = seat.c + dirC[j];

                if (rangeIn(r, c) && classroom[r][c] > 0) {
                    int tmp = classroom[r][c];

                    if (students.get(tmp).contains(number)) {
                        satisfaction[r][c]++;
                    }
                }
            }
        }

        // 2. 만족도 계산
        int sum = getSatisfaction();

        System.out.println(sum);
    }

    private static Seat seatInfo(int r, int c, Set likeFriends) {
        int empty = 0;
        int like = 0;

        for (int i = 0; i < 4; i++) {
            int nextR = r + dirR[i];
            int nextC = c + dirC[i];

            if (rangeIn(nextR, nextC)) {
                if (classroom[nextR][nextC] == 0) {
                    empty++;
                } else if (likeFriends.contains(classroom[nextR][nextC])) {
                    like++;
                }
            }
        }

        return new Seat(r, c, empty, like);
    }

    private static int getSatisfaction() {
        int sum = 0;

        for (int r = 0; r < N; r++) {
            for (int c = 0; c < N; c++) {
                switch (satisfaction[r][c]) {
                    case 1:
                        sum += 1;
                        break;
                    case 2:
                        sum += 10;
                        break;
                    case 3:
                        sum += 100;
                        break;
                    case 4:
                        sum += 1000;
                        break;
                }
            }
        }

        return sum;
    }

    private static boolean rangeIn(int r, int c) {
        return (r >= 0) && (c >= 0) && (r < N) && (c < N);
    }

    static class Seat implements Comparable<Seat> {
        int r, c, empty, like;

        public Seat(int r, int c, int empty, int like) {
            this.r = r;
            this.c = c;
            this.empty = empty;
            this.like = like;
        }

        @Override
        public int compareTo(Seat seat) {
            if (this.like == seat.like) {

                if (this.empty == seat.empty) {

                    if (this.r == seat.r) {
                        return this.c - seat.c;
                    }

                    return this.r - seat.r;
                }

                return seat.empty - this.empty;
            }

            return seat.like - this.like;
        }
    }
}