백준 14431번 소수마을 :: 마이구미

이번 글은 백준 알고리즘 14431번 "소수마을" 을 다뤄본다.

최근 선린고에서 주최된 머그컵이라는 대회에서 출제된 문제이다.

이전에 다뤘던 9205번 "맥주 마시면서 걸어가기" 와 흡사한 문제가 된다.

9205번 문제는 BFS도 좋은 접근이지만, 플로이드 알고리즘을 통해 문제 풀이를 했다. (9205번 풀이)

이번 문제에서는 플로이드가 아닌 BFS를 통해 문제를 해결해본다.

DFS, BFS - http://mygumi.tistory.com/102

Github 알고리즘 문제 - https://github.com/hotehrud/acmicpc

소수 마을들의 주민들은 매우 특이한 규칙을 준수한다. 규칙은 바로 “가고 싶은 위치까지의 거리가 소수일 경우에만 간다”라는 것이다. 소수 마을의 주민 승욱이는 소수 마을에서 멀리 떨어진 A마을에 볼일이 있어 그곳까지 가야한다. 소수 마을에서 A마을까지의 단숨에 가고 싶지만 안타깝게도 두 마을의 거리는 소수가 아닐 경우에는 그럴 수가 없다. 그럴 경우에는 다른 마을들을 경유하여 가야한다. (경유하는 마을도 현재 위치에서의 거리가 소수일 경우에만 갈 수 있다.) 소수 마을과 경유할 수 있는 마을들, 그리고 A마을의 위치가 좌표평면 상으로 주어질 때, 승욱이가 소수 마을의 규칙을 준수하여 A마을로 갈 수 있는 최단의 길을 찾는 것을 도와주자. 소수 판정을 위해 마을 간의 거리는 정수 부분만으로 취급한다. 예를 들어, 거리가 3.1415라면 이를 버림하여 3만 취급한다.

먼저 각 마을을 정점으로 잡고 인접행렬을 만들어준다.

거리가 소수일 경우에만 경유할 수 있기에 소수일 경우만을 고려한다.

for (int i = 0; i < n + 2; i++) {
    for (int j = 0; j < n + 2; j++) {
        if (i == j) {
            continue;
        }

        Point p1 = pos.get(i);
        Point p2 = pos.get(j);

        int distance = (int) Math.sqrt(((p2.x - p1.x) * (p2.x - p1.x)) + ((p2.y - p1.y) * (p2.y - p1.y)));

        if (isPrime(distance)) {
            adj[i].add(new Town(j, distance));
        }
    }
}

그 후 BFS를 돌린다.

무조건 소수 마을 위치에서 출발해야하기 때문에, 소수 마을과 연결된 정점들만 체크해야한다.

그렇기에 dist라는 배열을 만들어 모든 공간을 최대값으로 저장하여 최소 비용을 계산해야한다.

public static void bfs(ArrayList<Town>[] adj, int[] dist) {
    Queue<Integer> q = new LinkedList<>();

    q.add(0);
    dist[0] = 0;

    while (!q.isEmpty()) {
        int i = q.poll();

        for (Town town : adj[i]) {
            int j = town.idx;
            int d = town.distance;

            if (dist[i] != Integer.MAX_VALUE) {
                int cost = dist[i] + d;

                if (dist[j] > cost) {
                    dist[j] = cost;
                    q.add(j);
                }
            }
        }
    }
}

소수 체크하는 부분을 에라토스테네스의 체와 같은 알고리즘을 사용하면 조금 더 빠른 속도를 이끌어낼 수 있다.

전체 코드는 아래를 참고하길 바란다.

private void solve() {
    ArrayList<Town>[] adj = (ArrayList<Town>[]) new ArrayList[4002];
    ArrayList<Point> pos = new ArrayList<>();
    int[] dist = new int[4002];
    int startX = sc.nextInt();
    int startY = sc.nextInt();
    int endX = sc.nextInt();
    int endY = sc.nextInt();
    int n = sc.nextInt();
 
    for (int i = 0; i < n + 2; i++) {
        dist[i] = Integer.MAX_VALUE;
    }
 
    for (int i = 0; i < n + 2; i++) {
        adj[i] = new ArrayList<>();
    }
 
    pos.add(new Point(startX, startY));
    for (int i = 1; i <= n; i++) pos.add(new Point(sc.nextInt(), sc.nextInt()));
    pos.add(new Point(endX, endY));
 
    for (int i = 0; i < n + 2; i++) {
        for (int j = 0; j < n + 2; j++) {
            if (i == j) {
                continue;
            }
 
            Point p1 = pos.get(i);
            Point p2 = pos.get(j);
 
            int distance = (int) Math.sqrt(((p2.x - p1.x) * (p2.x - p1.x)) + ((p2.y - p1.y) * (p2.y - p1.y)));
 
            if (isPrime(distance)) {
                adj[i].add(new Town(j, distance));
            }
        }
    }
 
    bfs(adj, dist);
    if (dist[n + 1] == Integer.MAX_VALUE) {
        System.out.println(-1);
    } else {
        System.out.println(dist[n + 1]);
    }
}
 
public static boolean isPrime(int distance) {
    if (distance <= 2) {
        return false;
    }
 
    for (int i = 2; i <= Math.sqrt(distance); i++) {
        if (distance % i == 0) {
            return false;
        }
    }
    return true;
}
 
public static void bfs(ArrayList<Town>[] adj, int[] dist) {
    Queue<Integer> q = new LinkedList<>();
 
    q.add(0);
    dist[0] = 0;
 
    while (!q.isEmpty()) {
        int i = q.poll();
 
        for (Town town : adj[i]) {
            int j = town.idx;
            int d = town.distance;
 
            if (dist[i] != Integer.MAX_VALUE) {
                int cost = dist[i] + d;
 
                if (dist[j] > cost) {
                    dist[j] = cost;
                    q.add(j);
                }
            }
        }
    }
}
 
public static class Point {
    int x;
    int y;
 
    Point(int x, int y) {
        this.x = x;
        this.y = y;
    }
}
 
public static class Town {
    int idx;
    int distance;
 
    Town(int idx, int distance) {
        this.idx = idx;
        this.distance = distance;
    }
}