[Java] TreeSet에서 커스텀 객체 정렬

궁금증이 들어서 코드를 쳐보다가 알게 된 점이 있어서 TreeSet에 대해서 기록해두고자 합니다.

 

기본적으로 Set은 중복을 허용하지 않는 자료구조입니다. 여기에 Tree가 붙으면서 이진 탐색 트리(binary search tree)의 형태로 Set이 구현된 자료구조가 됩니다. 즉, 데이터들이 정렬이 되어 저장된다는 것입니다.

 

정렬이 된다는 부분에서 Integer 타입이나 String 타입이 들어갔을 때는 아 정렬이 쉽게 되겠구나라고 생각을 했습니다. 근데 만약 Student 같은 커스텀 객체가 들어간다면 어떻게 될까라는 궁금점이 생겨서 코드를 작성해 봤습니다.

 

# Student 객체 생성

Student 객체는 학번(studentId)과 이름(name)만을 필드로 갖는 객체입니다.

public class Student {
    private String name;
    private int studentId;


   public Student(int studentId, String name) {
       this.studentId = studentId;
       this.name = name;
   }

    public int getStudentId() {
        return studentId;
    }

    public String getName() {
        return name;
    }
}

 

 

# TreeSet에 Student 객체 넣기

TreeSet을 만들어서 이 안에 Student 객체를 넣어볼 것입니다.

public class RunStudent {
    public static void main(String[] args) {
        TreeSet<Student> school = new TreeSet<>();
        school.add(new Student(1, "김씨"));
    }
}

1이라는 학번과 "김씨"라는 이름을 갖는 학생을 하나 넣어줬습니다. 컴파일 오류도 안 나고 한번 실행시켜 봅니다.

Exception in thread "main" java.lang.ClassCastException: class comparator_comparable.Student cannot be cast to class java.lang.Comparable (comparator_comparable.Student is in unnamed module of loader 'app'; java.lang.Comparable is in module java.base of loader 'bootstrap') at java.base/java.util.TreeMap.compare(TreeMap.java:1291) at java.base/java.util.TreeMap.put(TreeMap.java:536) at java.base/java.util.TreeSet.add(TreeSet.java:255) at comparator_comparable.RunStudent.main(RunStudent.java:11)

ClassCastException이 발생합니다. 읽어보면 Student는 Comparable를 cast 하지 않았다는 어쩌고 쓰여있습니다. 이는 TreeSet은 데이터를 넣었을 때, 정렬을 해주는 자료구조이기 때문에 해당 객체의 정렬 기준이 필요합니다.

 

하지만 Student 클래스에는 학번과 이름이라는 2개의 필드가 존재하는데 Java에서 난 뭘로 정렬을 해야 할지 모르겠다는 예외를 발생시킨 겁니다. 그럼 Comparable 인터페이스를 사용해서 Student의 정렬 기준을 설정해 보겠습니다.

 

 

# Student implements Comparable

public class Student implements Comparable<Student> {
    private String name;
    private int studentId;

    ... 생략 ...
    
    @Override
    public int compareTo(Student std) {
       return this.studentId - std.studentId;
    }
}

Comparable 인터페이스를 implements 받으면 반드시 compareTo() 메서드를 오버라이드해서 재정의해줘야 합니다. 여기서 Student 클래스의 정렬 기준을 정의할 수 있습니다. 

 

compareTo() 메서드는 양수(1), 음수(-1), 0을 리턴값으로 사용할 수 있습니다. 위에서 작성한 compareTo()는 파라미터로 들어온 다른 Student 객체와 나 자신(this)을 비교합니다. 즉, 정렬의 기준이 studentId라는 것을 의미합니다.

 

나의 학번(this.studentId)과 다른 사람의 학번(std.studentId)를 뺀 값을 리턴합니다. 이 결괏값은 같으면 0, 내가 크면 양수, 쟤가 크면 음수의 리턴 값을 가집니다. 내가(this) 비교의 앞에 있기 때문에 위 코드는 학번을 기준으로 오름차순 정렬이 진행됩니다. 

 

그럼 내림차순 정렬을 하고 싶을 때는 어떻게 해야 할까?

@Override
public int compareTo(Student std) {
   return std.studentId - this.studentId;
}

이렇게 두 개를 바꿔서 리턴을 해주면 됩니다. Comparable의 compareTo()를 잘 이용해 주면 정렬 조건이 2개여도 쉽게 정렬을 할 수 있게 됩니다.

 

 

# TreeSet 다시 실행해 보기

이제 오류 없이 잘 돌아가는 것을 확인할 수 있습니다. 학생을 4명 정도 넣어서 정렬이 잘 되어서 나오는지 코드로 확인해 보겠습니다.

public class RunStudent {
    public static void main(String[] args) {
        TreeSet<Student> school = new TreeSet<>();
        school.add(new Student(1, "김씨"));
        school.add(new Student(2, "박씨"));
        school.add(new Student(3, "정씨"));
        school.add(new Student(4, "오씨"));

        for (Student student : school) {
            System.out.println("학번: " + student.getStudentId() + " 이름: " + student.getName());
        }
    }
}

<출력 결과>
학번: 1 이름: 김씨
학번: 2 이름: 박씨
학번: 3 이름: 정씨
학번: 4 이름: 오씨

의도한 대로 오름차순으로 정렬돼서 결과가 출력된 것을 확인할 수 있었습니다. 근데 한 가지 의문점은 여기서부터 시작됩니다. 만약 학번이 1인 학생이 한 명 더 추가될 때는 어떻게 될까? (같은 학번의 학생)

 

 

# 같은 학번을 가진 학생을 추가해 보기

public class RunStudent {
    public static void main(String[] args) {
        TreeSet<Student> school = new TreeSet<>();
        school.add(new Student(1, "김씨"));
        school.add(new Student(2, "박씨"));
        school.add(new Student(3, "정씨"));
        school.add(new Student(4, "오씨"));

        school.add(new Student(1, "이씨"));

        for (Student student : school) {
            System.out.println("학번: " + student.getStudentId() + " 이름: " + student.getName());
        }
    }
}

출력 결과를 보면 "이씨"는 출력되지 않았습니다. 즉 추가가 되지 않았다는 것을 볼 수 있습니다. 궁금해서 TreeSet의 contains() 메서드를 사용해서 어떤 결과가 나오는지 확인해 봤습니다.

school.add(new Student(1, "김씨"));
school.add(new Student(2, "박씨"));
school.add(new Student(3, "정씨"));
school.add(new Student(4, "오씨"));

Student lee = new Student(1, "이씨");

if (school.contains(lee)) {
    System.out.println(lee.getName() + "은 이미 학교에 존재합니다.");
} else {
    System.out.println(lee.getName() + "은 학교에 존재하지 않습니다.");
}

출력 결과는 "이씨은 이미 학교에 존재합니다."가 나왔습니다. school.contains()에 Integer나 String이 아닌 Student 객체를 넣어준 것인데, 존재 여부를 어떻게 판단할 수 있는 걸까 라는 의문이 들었습니다.

 

그래서 school에 존재하는 모든 학생 객체를 순회하면서 equals()를 사용하여 lee와 같은 객체가 있는지 확인해 봤습니다.

TreeSet<Student> school = new TreeSet<>();
school.add(new Student(1, "김씨"));
school.add(new Student(2, "박씨"));
school.add(new Student(3, "정씨"));
school.add(new Student(4, "오씨"));

Student lee = new Student(1, "이씨");

for (Student student : school) {
    if (student.equals(lee)) {
        System.out.println("lee와 같은 객체입니다.");
    }
}

결과는 아무것도 출력되지 않았습니다. 즉, "김씨", "박씨", "정씨", "오씨" 모두 "이씨"와는 다른 객체라는 것입니다. 근데 contains()에서 포함이 되어있다고 한다...?

 

궁금해서 Student에 다른 필드 값인 학년(grade)을 넣어서 확인해 봤습니다.

 

 

# Student 클래스에 grade 필드값 추가

public class Student implements Comparable<Student> {
    private String name;
    private int studentId;
    private int grade;


   public Student(int studentId, String name, int grade) {
       this.studentId = studentId;
       this.name = name;
       this.grade = grade;
   }

    ... 생략 ...

    public int getGrade() {
       return grade;
    }

    @Override
    public int compareTo(Student std) {
       return this.studentId - std.studentId;
    }
}

새롭게 추가해 주고, 다시 contains()를 사용해 봤습니다.

TreeSet<Student> school = new TreeSet<>();
school.add(new Student(1, "김씨", 1));
school.add(new Student(2, "박씨", 2));
school.add(new Student(3, "정씨", 3));
school.add(new Student(4, "오씨", 4));

Student lee = new Student(1, "이씨", 1);

System.out.println(school.contains(lee));

여전히 결괏값은 true가 나옵니다. lee가 school에 포함되어 있다는... school을 순회하면 모두 출력했을 때는 "이씨"가 존재하지 않습니다. 여기서 이제 "이씨"의 학번을 겹치지 않게 5로 바꾸고, 학년은 "김씨"와 겹치는 1로 유지한채 school에 넣어봤습니다.

TreeSet<Student> school = new TreeSet<>();
school.add(new Student(1, "김씨", 1));
school.add(new Student(2, "박씨", 2));
school.add(new Student(3, "정씨", 3));
school.add(new Student(4, "오씨", 4));

Student lee = new Student(5, "이씨", 1);
school.add(lee);

for (Student student : school) {
    System.out.println("학번: " + student.getStudentId() + " 이름: " + student.getName());
}

학번: 1 이름: 김씨
학번: 2 이름: 박씨
학번: 3 이름: 정씨
학번: 4 이름: 오씨
학번: 5 이름: 이씨

어..? grade는 겹치게 넣었는데, 이번에는 들어가네라는 의문에서 거의 확신이 섰습니다. 이진트리는 중복되는 숫자가 존재하지 않습니다. 즉, 우리가 Student를 만들 때 정렬 조건을 studentId로 지정했고 이것이 Student의 key 값이 되어 적용되고 있었습니다. 그래서 학번이 같으면 TreeSet에 넣어지지 않고, 학년이 같았을 때는 넣어졌던 것입니다.

 

이것에 확신을 얻기 위해서 Student의 정렬 조건을 학번(studentId)에서 학년(grade)으로 바꿔서 진행해 봤습니다.

@Override
public int compareTo(Student std) {
   return this.getGrade() - std.getGrade();
}

학번: 1 이름: 김씨
학번: 2 이름: 박씨
학번: 3 이름: 정씨
학번: 4 이름: 오씨

예상했던 대로 "김씨"와 학년이 같았던 "이씨"는 추가되지 않았습니다. 이렇게 TreeSet에 대해서... 좀 더 확실히 알 수 있는 계기가 되었던 거 같습니다.

 

자주 사용했던 자료구조는 아니었는데, 궁금해서 TreeSet을 사용한 출석부를 만들어봤습니다. 근데 그 과정에서 위와 같은 의문점이 들었고 TreeSet에 대해서 확실히 알게 되었던 과정이었습니다.. 👀