본문 바로가기
코딩테스트

[백준] 1303번 '전쟁 - 전투' - Java

by CuckooBird 2023. 7. 4.

문제


코드

맞았습니다가 뜬 코드입니다. - 메모리 212488KB | 시간 588ms | 코드 길이 3180B

import java.io.*;
import java.util.*;

public class Main {
	private static final BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));
	private static final BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
	
	private static int cnt = 0;
	private static boolean[] visited;
	
	public static void main(String[] args) throws IOException {
		StringTokenizer st = new StringTokenizer(bf.readLine());
		int n = Integer.parseInt(st.nextToken());
		int m = Integer.parseInt(st.nextToken());
		
		ArrayList<int[]> w_soldier = new ArrayList<>();
		ArrayList<int[]> b_soldier = new ArrayList<>();
		
		for(int i = 0 ; i < m ; i++) {
			String team = bf.readLine();
			for(int j = 0 ; j < n ; j++) {
				if('W' == team.charAt(j)) {
					int[] w = new int[2];
					w[0] = i;
					w[1] = j;
					w_soldier.add(w);
				} else {
					int[] b = new int[2];
					b[0] = i;
					b[1] = j;
					b_soldier.add(b);
				}
			}
		}
		
		int[][] w_node = draw_connection(w_soldier.size(), w_soldier);
		int[][] b_node = draw_connection(b_soldier.size(), b_soldier);
		
		// bfs 함수 사용 시 큐
		PriorityQueue<Integer> w_graph = new PriorityQueue<>();
		PriorityQueue<Integer> b_graph = new PriorityQueue<>();
		
		// dfs 함수 사용 시 스택
		//Stack<Integer> w_graph = new Stack<>();
		//Stack<Integer> b_graph = new Stack<>();
		
		visited = new boolean[w_soldier.size()];
		int w_sum = 0;
		for(int i = 0 ; i < w_soldier.size() ; i++) {
			cnt = 0;
			if(!visited[i]) {
				bfs(w_graph, w_node, i);
//				dfs(w_graph, w_node, i);
				w_sum += Math.pow(cnt, 2);
			}
		}
		
		visited = new boolean[b_soldier.size()];
		int b_sum = 0;
		for(int i = 0 ; i < b_soldier.size() ; i++) {
			cnt = 0;
			if(!visited[i]) {
				bfs(b_graph, b_node, i);
//				dfs(b_graph, b_node, i);
				b_sum += Math.pow(cnt, 2);
			}
		}
		
		bw.write(Integer.toString(w_sum));
		bw.write("\n");
		bw.write(Integer.toString(b_sum));
		
		bw.flush();
		bf.close();
		bw.close();
		
	}
	
	private static int[][] draw_connection(int size, ArrayList<int[]> arr) {
		int[][] path = new int[size][size];
		
		for(int i = 0 ; i < size ; i++) {
			for(int j = 0 ; j < i ; j++) {
				if(Math.abs(arr.get(j)[0] - arr.get(i)[0]) == 1 && arr.get(j)[1] - arr.get(i)[1] == 0) {
					path[i][j] = 1;
					path[j][i] = 1;
				} else if(Math.abs(arr.get(j)[1] - arr.get(i)[1]) == 1 && arr.get(j)[0] - arr.get(i)[0] == 0) {
					path[i][j] = 1;
					path[j][i] = 1;
				}
			}
		}
		
		return path;
	}
	
	private static void bfs(PriorityQueue<Integer> graph, int[][] soldier_node, int node) {
		visited[node] = true;
		graph.poll();
		cnt++;
		for(int i = 0 ; i < soldier_node.length ; i++) {
			if(!visited[i] && soldier_node[node][i] == 1) {
				graph.add(i);
				bfs(graph, soldier_node, i);
			}
		}
	}
	
	private static void dfs(Stack<Integer> graph, int[][] soldier_node, int node) {
		visited[node] = true;
		graph.push(node);
		cnt++;
		for(int i = 0 ; i < soldier_node.length ; i++) {
			if(!visited[i] && soldier_node[node][i] == 1) {
				dfs(graph, soldier_node, i);
			} else if(visited[i] && soldier_node[node][i] == 1 && graph.size() > 0) {
				graph.pop();
			}
		}
	}
}

 

1. 노드 간 간선 만들기 - draw_connection 함수

private static int[][] draw_connection(int size, ArrayList<int[]> arr) {
		int[][] path = new int[size][size];
		
		for(int i = 0 ; i < size ; i++) {
			for(int j = 0 ; j < i ; j++) {
				if(Math.abs(arr.get(j)[0] - arr.get(i)[0]) == 1 && arr.get(j)[1] - arr.get(i)[1] == 0) {
					path[i][j] = 1;
					path[j][i] = 1;
				} else if(Math.abs(arr.get(j)[1] - arr.get(i)[1]) == 1 && arr.get(j)[0] - arr.get(i)[0] == 0) {
					path[i][j] = 1;
					path[j][i] = 1;
				}
			}
		}
		
		return path;
	}

ArrayList<int[]> arr는 좌표가 배열형태로 들어가있는 리스트입니다. 만약 W를 (2, 3) 에 두었다고 한다면 [2, 3] 가 arr의 요소입니다.

path는 노드를 모두 연결한 뒤에 리턴할 2차원 배열입니다. 0이라면 연결이 되어있지 않고, 1이라면 연결이 된 상태입니다.

그리고 연결은 상, 하, 좌, 우 로 붙어있는지 확인하여 연결시킵니다.

 

2. 넓이 우선 탐색 함수 - bfs 함수

private static void bfs(PriorityQueue<Integer> graph, int[][] soldier_node, int node) {
		visited[node] = true;
		graph.poll();
		cnt++;
		for(int i = 0 ; i < soldier_node.length ; i++) {
			if(!visited[i] && soldier_node[node][i] == 1) {
				graph.add(i);
				bfs(graph, soldier_node, i);
			}
		}
	}

bfs는 큐를 이용해서 연결된 노드들을 순서대로넣고, 탐색이 끝나는대로 빼서 다시 연결된 노드들을 넣는 방식으로 진행했습니다.

 

3. 깊이 우선 탐색 함수 - dfs 함수

private static void dfs(Stack<Integer> graph, int[][] soldier_node, int node) {
		visited[node] = true;
		graph.push(node);
		cnt++;
		for(int i = 0 ; i < soldier_node.length ; i++) {
			if(!visited[i] && soldier_node[node][i] == 1) {
				dfs(graph, soldier_node, i);
			} else if(visited[i] && soldier_node[node][i] == 1 && graph.size() > 0) {
				graph.pop();
			}
		}
	}

dfs는 스택을 이용하였습니다. 연결된 노드를 따라 쭉쭉 스택에 집어넣고, 연결된 부분이 없다면 스택을 pop하여 다시 돌아가는 방식으로 진행했습니다.

 

4. 메인함수에서 bfs, dfs 함수 이용하기

// bfs 함수 사용 시 큐
		PriorityQueue<Integer> w_graph = new PriorityQueue<>();
		PriorityQueue<Integer> b_graph = new PriorityQueue<>();
		
		// dfs 함수 사용 시 스택
		//Stack<Integer> w_graph = new Stack<>();
		//Stack<Integer> b_graph = new Stack<>();
		
		visited = new boolean[w_soldier.size()];
		int w_sum = 0;
		for(int i = 0 ; i < w_soldier.size() ; i++) {
			cnt = 0;
			if(!visited[i]) {
				bfs(w_graph, w_node, i);
//				dfs(w_graph, w_node, i);
				w_sum += Math.pow(cnt, 2);
			}
		}
		
		visited = new boolean[b_soldier.size()];
		int b_sum = 0;
		for(int i = 0 ; i < b_soldier.size() ; i++) {
			cnt = 0;
			if(!visited[i]) {
				bfs(b_graph, b_node, i);
//				dfs(b_graph, b_node, i);
				b_sum += Math.pow(cnt, 2);
			}
		}

bfs 함수를 이용하려면 위 코드를 그대로 실행시키면 되고, dfs 함수를 이용하려면 위 코드에서 주석 된 부분을 다시 실행시키고 큐와 bfs 함수 사용 부분을 주석시키면 됩니다.


후기

어제 쓴 코드 거의 다 가져왔는데 W와 B로 두 가지 경우여서 정적 변수를 사용하는 걸 자제한 게 차이입니다. 이게 될까 라는 생각을 했는데 이게 되네 해서 조금은 놀랐습니다. 예전에 함수는 C언어에서 많이 데여서 함수를 쓰는 걸 좀 꺼려했는데 역시 자바 최고..

계속 런타임 에러 (StringIndexOutOfBounds) 가 떠서 왜이런가 했더니 가로의 갯수가 n개, 세로의 갯수가 m개라서 nxm 행렬이 아니라 mxn 행렬이더라고요.. 허망함..