[Java] TreeSet 사용법

사실 사용법이라기 보다는, TreeSet을 사용할 때 유의할 점 정도가 맞겠다.

HashSet을 사용하기 위해서는 HashSet에 들어가는 원소의 equals() 함수와 hashCode() 함수를 각각 override해줘야 한다. 그럼 TreeSet을 사용하기 위해서는?

Oracle이 제공하는 Java Doc 문서를 살펴보면, 다음과 같이 나와있다.
This is so because the Set interface is defined in terms of the equals operation, but a TreeSet instance performs all element comparisons using its compareTo (or compare) method, so two elements that are deemed equal by this method are, from the standpoint of the set, equal. 
==> Set 인터페이스는 equals() 함수를 통하여 정의되어 있는데, TreeSet은 모든 원소의 비교를 compareTo() 혹은 compare() 함수를 통해서 수행하기 때문에, set의 관점에서 봤을 때, 두 원소들이 compareTo() 함수에 의해 같다고 생각되면, 동일해야 한다.
출처 : 
http://download.oracle.com/javase/6/docs/api/java/util/TreeSet.html
즉, TreeSet은 원소의 compareTo() 함수만 잘 구현해주면, 정렬된 상태로 잘 저장이 된다는 말씀! 실제로 TreeSet 예제를 찾아보면, compareTo() 함수만 override 한 예제도 많다. 그리고 equals()가 참일 때, compareTo() 함수가 0을 반환하도록만 신경 써 주면 별 문제없이 잘 동작한다.

그럼 이제 주의해야 할 점! 아래 코드를 보자.
import java.util.Set;
import java.util.TreeSet;


public class Test {
	public static class Data implements Comparable<Data>{
		private String id = null;
		private int priority = 0;
		public Data(String id, int priority){
			this.id = id;
			this.priority = priority;
		}
		@Override
		public int compareTo(Data op){	
			System.out.printf("%s vs %s\n", this, op);
			// id가 동일하면 동일한 객체
			if( id.equals( op.id ) ){
				return 0;
			}			
			if( priority != op.priority ){
				if( priority > op.priority ){
					return 1;
				}
				else{
					return -1;
				}
			}
			return id.compareTo( op.id );			
		}
		@Override
		public String toString(){
			return String.format("%s/%d", id, priority);
		}	
		
	}
	public static void main(String[] args){
		Set<Data> treeSet = new TreeSet<Data>();		
		
		treeSet.add( new Data("a", 0) );
		treeSet.add( new Data("b", 1) );
		treeSet.add( new Data("c", 2) );
		treeSet.add( new Data("a", 3) );
		
		for(Data d : treeSet ){
			System.out.println( d );
		}		
	}

}
위 코드는 심플하다. TreeSet에 들어가는 Data class는 String id를 통해 객체의 동일 여부를 판별한다. 그리고 id가 다르면, int priority 변수를 통해 정렬된다. 이 코드를 실행 결과는 어떻게 될까? 상식적으로 판단하면, 마지막에 추가된 a/3 객체는 a/1 객체와 동일한 객체이기 때문에 TreeSet에 추가되면 안된다. 하지만 결과는 다음과 같다.

== 실행 결과==
a/0
b/1
c/2
a/3
==========

a/3가 추가되었으므로, 원하는 결과와는 다르다. equal()와 hashCode()를 추가해줘야 할까? 다음과 같이 두 함수를 추가해보자.
import java.util.Set;
import java.util.TreeSet;

public class Test {
	public static class Data implements Comparable<Data>{
		private String id = null;
		private int priority = 0;
		public Data(String id, int priority){
			this.id = id;
			this.priority = priority;
		}
		@Override
		public int compareTo(Data op){	
			System.out.printf("%s vs %s\n", this, op);
			// id가 동일하면 동일한 객체
			if( id.equals( op.id ) ){
				return 0;
			}			
			if( priority != op.priority ){
				if( priority > op.priority ){
					return 1;
				}
				else{
					return -1;
				}
			}
			return id.compareTo( op.id );			
		}
		@Override
		public String toString(){
			return String.format("%s/%d", id, priority);
		}	
		@Override
		public int hashCode(){
			return id.hashCode();
		}
		@Override
		public boolean equals(Object o){
			if( o == null)
				return false;
			if( !( o instanceof Data) ){
				return false;
			}
			Data op = (Data) o;
			if( this == op ){
				return true;
			}
			return id.equals( op.id );
		}
	}
	public static void main(String[] args){
		Set<Data> treeSet = new TreeSet<Data>();		
		treeSet.add( new Data("a", 0) );
		treeSet.add( new Data("b", 1) );
		treeSet.add( new Data("c", 2) );
		treeSet.add( new Data("a", 3) );
		
		for(Data d : treeSet ){
			System.out.println( d );
		}		
	}

}

하지만 여전히 실행결과는 다음과 같다.

== 실행 결과==
a/0
b/1
c/2
a/3
==========

무엇이 문제일까? 추측컨데, TreeSet은 내부적으로 tree 형태로 데이터를 저장할 것이다. a/0, b/1, c/2 객체가 추가되었을 때, well balanced binary tree 형태로 생각해보면, b/1가 tree의 root 이고, a/0와 c/2는 각각의 child node일 것이다. 여기에 새로운 a/3가 추가되었을 때, 이 a의 priority는 3이므로, 처음에 b/1과 compareTo() 함수가 실행되고, 그리고 다음으로는 c/2와 compareTo() 함수가 실행될 것이다. (실제 compareTo() 에 로그를 남기면 추측이 맞음을 알 수 있다.)

즉, 마지막에 추가된 a/3 객체는 a/0 객체와는 아예 비교가 일어나지 않으므로, 문제가 발생한다. 의미적으로 보면 동일한 두 객체인 a/3 객체와 a/0 객체가 모두 Set에 존재하는 상황이 발생하는 것이다.

그렇다면 어떻게 해야할까?

일단 나는 Java 내공이 부족한 관계로, TreeSet 외에 HashSet을 추가하여, TreeSet에 Data를 넣고 빼기 전에 HashSet을 통해 검증하는 방법을 택했다. 이게 맞는 방법인지는 모르겠지만, 일단은 제일 무난한 방법. 혹시 다른 해법이 있으신 분은 제보 바랍니다!




댓글

Designed by JB FACTORY