코딩 테스트(JAVA)/인프런 문제풀이

친구인가? (Disjoint-Set : Union&Find)

Lea Hwang 2022. 3. 29. 00:36

설명

오늘은 새 학기 새로운 반에서 처음 시작하는 날이다. 현수네 반 학생은 N명이다. 현수는 각 학생들의 친구관계를 알고 싶다.

모든 학생은 1부터 N까지 번호가 부여되어 있고, 현수에게는 각각 두 명의 학생은 친구 관계가 번호로 표현된 숫자쌍이 주어진다.

만약 (1, 2), (2, 3), (3, 4)의 숫자쌍이 주어지면 1번 학생과 2번 학생이 친구이고, 2번 학생과 3번 학생이 친구, 3번 학생과 4번 학생이 친구이다.

그리고 1번 학생과 4번 학생은 2번과 3번을 통해서 친구관계가 된다.

학생의 친구관계를 나타내는 숫자쌍이 주어지면 특정 두 명이 친구인지를 판별하는 프로그램을 작성하세요.

두 학생이 친구이면 “YES"이고, 아니면 ”NO"를 출력한다.

 

입력

첫 번째 줄에 반 학생수인 자연수 N(1<=N<=1,000)과 숫자쌍의 개수인 M(1<=M<=3,000)이 주어지고,

다음 M개의 줄에 걸쳐 숫자쌍이 주어진다.

마지막 줄에는 두 학생이 친구인지 확인하는 숫자쌍이 주어진다.

 

출력

첫 번째 줄에 “YES"또는 "NO"를 출력한다.

예시 입력 1 

9 7
1 2
2 3
3 4
1 5
6 7
7 8
8 9
3 8

예시 출력 1

NO

 

문제 이해

입력값을 토대로 이해한 내용을 그림으로 이해 후 수도 코드를 작성해보겠습니다.

Disjoint-Set : Union&Find를 이용하는 것 같지만, 아직 학습 전 이므로 우선 알고 있는 것들을 이용해서 최대한 코드를 작성하겠습니다. 

 

[코딩 테스트(JAVA)] - 회의실 배정

tmp 활용해서 코드 작성 중...

package ch09;

import java.util.*;

class Student implements Comparable<Student> {
	public int stu1, stu2;
	
	// 매개변수 있는 생성자
	Student(int stu1, int stu2) {
		this.stu1 = stu1;
		this.stu2 = stu2;
	}
	
	 public int arrInfo() {
			return stu2;
	 }
	
	// 오버라이드 compareTo를 통한 오름차순 정렬
	@Override
	public int compareTo(Student st) {
		// stu1 기준으로 오름차순 정렬
		// stu1가 같다면 stu2로 오름차순 정렬
		if(this.stu1 != st.stu1) {
			return this.stu1 - st.stu1;
		}else return this.stu2 - st.stu2;
	}
}

public class Friends {
	static ArrayList<Student> arr;
	
	public String solution(int friends1, int friends2) {
		String answer = "";
		// 정렬, arr 사용해야하므로 위에 static 선언
		Collections.sort(arr);
		// foreach문 돌면서 비교
		int tmp = Integer.MIN_VALUE;
		for(Student ob : arr) {
			//System.out.println(ob.arrInfo()); // stu2값 쭉 나옴 
			if(ob.stu1 > tmp) tmp = ob.stu2;
			
			// 채점결과 :  return 값 안나옴
			if(tmp == ob.stu1) {
				answer = "YES";
			}else answer ="NO";
		}
		return answer;
	}

	public static void main(String[] args) {
		Friends T = new Friends();
		// 입력 3개 : 학생 수(n), 숫자쌍의 개수(m), 숫자쌍
		Scanner kb = new Scanner(System.in);
		int n = kb.nextInt();
		int m = kb.nextInt();
		// 숫자쌍을 담을 배열 생성
		arr = new ArrayList<>();
		for(int i=0; i<=m; i++) {
			// m만큼 돌면서 순서쌍 넣기
			int stu1 = kb.nextInt();
			int stu2 = kb.nextInt();
			arr.add(new Student(stu1, stu2));
			// 마지막 입력 쌍은 질문
			if(i==m) {
				T.solution(stu1, stu2);
			}
		}
		
	}

}

채점 결과를 받아봤을 때 return값 자체가 안 나오는 에러가 발생했습니다.

 


이쯤 해서 강의를 들어보면서 수정 보완해보도록 하겠습니다.

 

선생님께서는 서로소집합을 통해서 문제를 해결하셨고 그때 Union&Find알고리즘을 통해 구현하셨습니다.

글보다는 그림으로 쉽게 설명해보도록하겠습니다.

 

코드를 보면서 입력 예시를 이해하는 게 빠르므로 코드를 보면서 진행하겠습니다.

 

💡 Find( ), Union( ) 코드는 암기합니다!

import java.util.*;
class Main {
	static int[] arr;
	public static int Find(int v){
		if(v==arr[v]) return v;
		else return arr[v]=Find(arr[v]);
	}
	public static void Union(int a, int b){
		int fa=Find(a);
		int fb=Find(b);
		if(fa!=fb) arr[fa]=fb;
	}
	public static void main(String[] args){
		.......
	}
}

 

이제 완전한 정답 코드를 작성해보겠습니다.

package ch09;

import java.util.*;

public class Friends {
	static int[] arr;
	
	public static int Find(int v) {
		// 집합번호 리턴하는 메서드
		if(v == arr[v]) return v;
		else return arr[v] = Find(arr[v]);
	}
	
	public static void Union(int a, int b) {
		// 둘을 한 집합으로 만드는 메서드
		int fa = Find(a);
		int fb = Find(b);
		if(fa != fb) arr[fa] = fb;
	}
	
	public static void main(String[] args) {
		Scanner kb = new Scanner(System.in);
		int n = kb.nextInt();
		int m = kb.nextInt();
		// 초기화 배열 생성 : 학생번호와 집합번호 일치
		arr = new  int[n+1]; // n번학생까지 인덱스가 나와야 하므로 n+1
		for(int i=1; i<=n; i++) arr[i] = i;
		// 순서쌍 받음
		for(int i=1; i<=m; i++ ) {
			int a = kb.nextInt();
			int b = kb.nextInt();
			Union(a,b);
		}
		// 입력예시 출력
		int a = kb.nextInt();
		int b = kb.nextInt();
		int fa = Find(a);
		int fb = Find(b);
		if(fa == fb) System.out.println("YES");
		else System.out.println("NO");
		
	}

}

 

💡 arr배열에서 인덱스가 학생수 n만큼 나와야 하므로

arr = new int [n]이 아니라 arr = new int [n+1]로 작성하였습니다.

 

 


혼자 해왔던 것을 쥐어짜내서 코드 작성(5시간) 후 강의를 보고 익힌 후 코드 작성(5시간)하니까 이틀이 훅 갔습니다.

강하게 현타가 오지만 나중에는 척척 풀어낼 수 있겠죠...?

 

그래도 다행인 점은 10시간 동안 한 문제를 풀 때 힘들기는 했지만 재미가 없다든지 하기 싫다는 생각은 안 들었던 것 같습니다. 

 

앞으로도 천천히 그리고 확실히 공부해서 코딩 테스트 문제 중간 레벨 정도는 풀 수 있는 날이 하루빨리 왔으면 합니다!