尋常でないもふもふ

a software engineer blog

TreeSet の等価性判定は equals() ではない件

下記は 1 ユーザーのスコアを保持する Entity。
これを 100 個作って TreeSet に突っ込んだら、100 個なかったという話。

import java.util.Random;
import java.util.Set;
import java.util.TreeSet;

public class UserScore implements Comparable<UserScore> {

    // ユーザーID
    private int userId;
    // スコア
    private int score;

    public int getUserId() {
        return userId;
    }
    public void setUserId(int userId) {
        this.userId = userId;
    }
    public int getScore() {
        return score;
    }
    public void setScore(int score) {
        this.score = score;
    }

    @Override
    public int hashCode() {
        final int prime = 31;
        int result = 1;
        result = prime * result + userId;
        return result;
    }
    @Override
    public boolean equals(Object obj) {
        if (this == obj)
            return true;
        if (obj == null)
            return false;
        if (!(obj instanceof UserScore))
            return false;
        UserScore other = (UserScore) obj;
        if (userId != other.userId)
            return false;
        return true;
    }

    @Override
    public int compareTo(UserScore o) {
        return score - o.getScore();
    }

    // テストコード
    public static void main(String[] args) {
        Random r = new Random();
        Set<UserScore> set = new TreeSet<>();
        
        // 100個の UserScore を作成
        for (int i = 1; i <= 100; i++) {
            UserScore entity = new UserScore();
            entity.setUserId(i);
            entity.setScore(r.nextInt(11)); //0点~10点のランダム値
            set.add(entity);
        }
        
        // すべて出力
        for (UserScore score : set) {
            System.out.printf("user[%2d] %2dpt%n", score.getUserId(), score.getScore());
        }
        System.out.printf("size = %d%n", set.size());
    }


出力結果

user[20]  0pt
user[30]  1pt
user[12]  2pt
user[ 3]  3pt
user[50]  4pt
user[ 1]  5pt
user[ 2]  6pt
user[ 8]  7pt
user[ 5]  8pt
user[ 4]  9pt
user[ 9] 10pt
size = 11


メインメソッドを実行すると、100人分作ったはずなのに、11個しかない。
equals() 実装したから、userId 基準でユニークになると思ったけど、TreeSet は equals() をみないそうだ。
equals() は HashSet 用にそのまま残しておいて、別途 compareTo() で等価かどうかを判定する処理を書く。

    @Override
    public int compareTo(UserScore o) {
        if (score < o.getScore()) {
            return -1;
        } else if (o.getScore() < score) {
            return 1;
        } else {
            // userId で等価性を判定
            if (userId < o.getUserId()) {
                return -1;
            } else if (o.getUserId() < userId) {
                return 1;
            } else {
                return 0;
            }
        }
    }

TreeSet の Javadoc みたら前文に書いてた。

TreeSet インスタンスはその compareTo または compare メソッドを使用してすべての要素比較を実行するので、このメソッドによって等価と見なされる 2 つの要素は、このセットの見地からすれば同じものです。