(BOJ) 핑크 플로이드_6091번

문제

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

접근

한참 고민하다가 정점 A에서 정점 B로 가는 비용이 작다는 것은 인접해있을 확률이 높다고 생각하고 비용 순서대로 우선순위 큐를 정의하였다.

그리고 우선순위 큐에서 하나씩 간선에 대한 정보를 뽑고 현재까지 만들어진 그래프를 탐색하면서 존재하지 않은 정점이면 붙이는 식으로 구현했었다.

코드 1 (메모리 초과)

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.PriorityQueue;

public class Main {

  static class Node implements Comparable<Node> {

    int x;
    int y;
    int value;

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

    @Override
    public int compareTo(Node n) {
      if (this.value > n.value) {
        return 1;
      } else if (this.value == n.value) {
        if (this.x > n.x) {
          return 1;
        } else if (this.x == n.x) {
          if (this.y > n.y) {
            return 1;
          }
        }
      }
      return -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 n = Integer.parseInt(br.readLine());
    int[][] map = new int[n + 1][n + 1];
    PriorityQueue<Node> pq = new PriorityQueue<>();

    for (int i = 1; i < n; i++) {
      String[] inputs = br.readLine().split(" ");
      for (int j = i + 1; j <= n; j++) {
        map[i][j] = Integer.parseInt(inputs[j - (i + 1)]);
        pq.add(new Node(j, i, map[i][j]));
      }
    }

    List<Integer>[] graph = new ArrayList[n + 1];
    for (int i = 1; i < graph.length; i++) {
      graph[i] = new ArrayList<>();
    }

    // x, y 순으로 넣을 예정
    PriorityQueue<Integer> posQ = new PriorityQueue<>();
    boolean[] visited = new boolean[n + 1];

    while (!pq.isEmpty()) {
      Node node = pq.remove();

      for (int i = 1; i <= n; i++) {
        visited[i] = false;
      }

      // 출발점 y, 도착점 x
      int x = node.x;
      int y = node.y;
      posQ.add(y);
      visited[y] = true;
      boolean hasAnswer = true;

      while (!posQ.isEmpty()) {
        int pos = posQ.remove();

        for (int i = 0; i < graph[pos].size(); i++) {
          int val = graph[pos].get(i);
          if (val == x) {
            hasAnswer = false;
            break;
          }

          if (!visited[val]) {
            posQ.add(val);
            visited[val] = true;
          }
        }
      }

      if (hasAnswer) {
        graph[y].add(x);
        graph[x].add(y);
      }

      posQ.clear();
    }

    for (int i = 1; i <= n; i++) {
      int size = graph[i].size();
      bw.write(size + " ");
      Collections.sort(graph[i]);
      for (int j = 0; j < size; j++) {
        bw.write(graph[i].get(j) + " ");
      }
      bw.write("\n");
    }

    bw.flush();
    bw.close();
    br.close();
  }
}

하지만 그 결과 메모리 초과가 나버렸다...
Input값이 커졌을 때, 도착점과 같은 정점이 나타나지 않았을 경우 중간에 큐에 정점에 대한 값이 쌓이게 되고, 이 과정에서 메모리 초과가 난 것으로 추정된다.

(그리고 다시 보니 우선순위 큐로 둘 필요도 없다.)

계속 고민하다가 한 가지 사실을 깨닫게 되었다. 생성된 그래프를 직접 탐색하는 대신, 두 정점이 현재 연결된 상태인지만 확인하면 끝나는 문제였다.

즉, 두 정점이 같은 그룹에 속해있는지만 확인하면 된다. Prim 알고리즘과 유사하다!

코드 2

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.PriorityQueue;

public class Main {

  static class Node implements Comparable<Node> {

    int x;
    int y;
    int value;

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

    @Override
    public int compareTo(Node n) {
      if (this.value > n.value) {
        return 1;
      } else if (this.value == n.value) {
        if (this.x > n.x) {
          return 1;
        } else if (this.x == n.x) {
          if (this.y > n.y) {
            return 1;
          }
        }
      }
      return -1;
    }
  }

  static int[] group;

  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 n = Integer.parseInt(br.readLine());
    int[][] map = new int[n][n];
    group = new int[n];
    PriorityQueue<Node> pq = new PriorityQueue<>();

    for (int i = 0; i < n; i++) {
      group[i] = i;
    }

    for (int i = 0; i < n - 1; i++) {
      String[] inputs = br.readLine().split(" ");
      for (int j = i + 1; j < n; j++) {
        map[i][j] = Integer.parseInt(inputs[j - (i + 1)]);
        pq.add(new Node(j + 1, i + 1, map[i][j]));
      }
    }

    List<Integer>[] graph = new ArrayList[n + 1];
    for (int i = 1; i < graph.length; i++) {
      graph[i] = new ArrayList<>();
    }

    while (!pq.isEmpty()) {
      Node node = pq.remove();

      // 출발점 y, 도착점 x
      int x = node.x;
      int y = node.y;

      if (find(x - 1) != find(y - 1)) {
        graph[y].add(x);
        graph[x].add(y);
        union(x - 1, y - 1);
      }
    }

    for (int i = 1; i <= n; i++) {
      int size = graph[i].size();
      bw.write(size + " ");
      Collections.sort(graph[i]);
      for (int j = 0; j < size; j++) {
        bw.write(graph[i].get(j) + " ");
      }
      bw.write("\n");
    }

    bw.flush();
    bw.close();
    br.close();
  }

  static void union(int a, int b) {
    int ap = find(a);
    int bp = find(b);

    group[ap] = bp;
  }

  static int find(int n) {
    if (n == group[n]) {
      return n;
    }

    group[n] = find(group[n]);
    return group[n];
  }
}

플로이드 와샬을 공부하다가 접하게 되었다. 드디어 첫 플래티넘 문제를 풀어보았다!