P.16946 벽 부수고 이동하기 4

코드

import java.io.*;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.LinkedList;

public class Main {

    static int n, m;
    static int[][] map;
    static int[][] copyMap;
    static ArrayList<Loc> walls = new ArrayList<>();
    static int[] idxArray;

    static public class Loc {
        int x, y;

        public Loc(int x, int y) {
            this.x = x;
            this.y = y;
        }
    }

    static int[] dx = {-1, 0, 1, 0};
    static int[] dy = {0, -1, 0, 1};

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

        int[] line = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
        n = line[0]; m = line[1];
        map = new int[n][m];
        copyMap = new int[n][m];
        for (int i = 0; i < n; i++) {
            map[i] = Arrays.stream(br.readLine().split("")).mapToInt(Integer::parseInt).toArray();
            for (int j = 0; j < m; j++) {
                copyMap[i][j] = map[i][j];
                if (map[i][j] == 1) {
                    walls.add(new Loc(j, i));
                    copyMap[i][j] = -1;
                }
            }
        }
        idxArray = new int[n * m + 1];

        mapCheck();

        for (Loc wall : walls) {
            int cnt = 1;

            HashSet<Integer> idxes = new HashSet<>();

            for (int i = 0; i < 4; i++) {
                int nextY = wall.y + dy[i];
                int nextX = wall.x + dx[i];

                if (nextY >= 0 && nextY < n && nextX >= 0 && nextX < m && copyMap[nextY][nextX] != -1)
                    idxes.add(copyMap[nextY][nextX]);
            }

            for (Integer idx : idxes) {
                cnt += idxArray[idx];
            }

            map[wall.y][wall.x] = cnt % 10;
        }

        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) bw.write(map[i][j] + "");
            bw.newLine();
        }
        bw.flush();
    }

    private static void mapCheck() {
        boolean[][] visited = new boolean[n][m];

        int idx = 1;
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < m; j++) {
                if (!visited[i][j] && map[i][j] == 0) {
                    bfs(visited, i, j, idx++);
                }
            }
        }
    }

    private static void bfs(boolean[][] visited, int y, int x, int idx) {
        LinkedList<Loc> queue = new LinkedList<>();
        queue.add(new Loc(x, y));

        int cnt = 0;
        while (!queue.isEmpty()) {
            Loc poll = queue.poll();

            if (!visited[poll.y][poll.x]) {
                visited[poll.y][poll.x] = true;
                copyMap[poll.y][poll.x] = idx;
                cnt++;

                for (int i = 0; i < 4; i++) {
                    int nextY = poll.y + dy[i];
                    int nextX = poll.x + dx[i];

                    if (nextY >= 0 && nextY < n && nextX >= 0 && nextX < m && !visited[nextY][nextX] && copyMap[nextY][nextX] != -1) {
                        queue.add(new Loc(nextX, nextY));
                    }
                }
            }
        }

        idxArray[idx] = cnt;
    }
}

코드 설명

코드가 조금 길지만 생각보다 알고리즘은 간단하다.

 

일단 벽마다 bfs를 돌면 무조건 시간이 초과된다.

따라서 벽 하나에 상, 하, 좌, 우 칸을 한 번씩만 확인해도 되는 코드를 작성했다.

 

벽을 뚫고 이동할 수 있는 0의 공간을 하나의 그룹으로 보았다.

그 그룹의 인덱스를 정하고 그룹의 요소의 개수를 인덱스와 함께 저장해 두었다.

다음의 예시일 때 그룹을 저장한 형식은 다음과 같다.

4 5
11001
00111
01010
10101

벽	벽	1(2)	1(2)	벽
2(3)	2(3)	벽	벽	벽
2(3)	벽	3(1)	벽	4(1)
벽	5(1)	벽	6(1)	벽

괄호밖의 숫자는 인덱스(그룹 번호)를 의미하고 괄호 안의 숫자는 그 그룹의 요소의 개수이다.

 

이제 벽의 입장에서는 상, 하, 좌, 우 칸을 한 번씩만 확인해 인덱스를 가져와 계산해주면 된다.

이때 (2, 1) 칸의 벽과 같이 상, 좌 칸의 인덱스가 같은 경우 두 번 더해지는 것을 피해야한다.

따라서 상, 하, 좌, 우의 인덱스를 집합으로 저장해서 중복을 피했다.

HashSet<Integer> idxes = new HashSet<>();

for (int i = 0; i < 4; i++) {
    int nextY = wall.y + dy[i];
    int nextX = wall.x + dx[i];

    if (nextY >= 0 && nextY < n && nextX >= 0 && nextX < m && copyMap[nextY][nextX] != -1)
        idxes.add(copyMap[nextY][nextX]);
}

for (Integer idx : idxes) {
    cnt += idxArray[idx];
}

map[wall.y][wall.x] = cnt % 10;