Problem when adding elements to a TreeSet – some elements are not added

3 min read >

Problem when adding elements to a TreeSet – some elements are not added

Engineering Insights & Enterprise solutions

I thought of sharing with you this little problem that I encountered when using a TreeSet to display an ordered list of items:

I have some items stored in the database, each of them having a unique id.
Some of the items have the same values, only the ids are different.
I also have a business method that creates a list of items, sorted by some custom ordering of fields.
To sort them, I used the following lines of code:

SortedSet result = new TreeSet(itemsComparator);
for (Item item : someItems) {
    if (itemIsEligible(item)) {
        result.add(item);
    }
}

where itemsComparator is an instance of this class:

public class ItemsComparator implements Comparator {
    /**
     * Compare two items for order.
     * @param item1 the first item
     * @param item2 the second item
     * @return -1 if item1 should appear before item2,
               1 if item2 should appear before item1,
               or 0 if the order is not important
     */
    public int compare(Item item1, Item item2) {
        // Sort by field1 descending
        int compareField1 = item1.getField1().compareTo(item2.getField1());
        if (compareField1 != 0) {
            return -compareField1;
        }
        // For equal field1, sort by field2 ascending
        int compareField2 = item1.getField2().compareTo(item2.getField2());
        return compareField2;
    }
}

At some moment I discovered the following bug: if two items had the same values, but different ids, only the last one of them appeared in the result.</>

After a little debugging I discovered that the problem was this:
TreeSet uses internally a TreeMap to store the elements.
TreeMap uses compareTo to determine the order of the elements, to know where to place them in the internal tree.
And if two elements are different but compareTo returns 0, then TreeMap will consider them equal, and the newly added element replaces the old one.

My solution to this problem was the following:

public class ItemsComparator implements Comparator {
    /**
     * Compare two items for order.
     * To ensure that two different items are not treated as equal,
     * our method will never return 0 for different items!
     * @param item1 the first item
     * @param item2 the second item
     * @return -1 if item1 should appear before item2,
                    1 if item2 should appear before item1,
                    0 if item1 and item2 are the same
     */
    public int compare(Item item1, Item item2) {
        // Sort by field1 descending
        int compareField1 = item1.getField1().compareTo(item2.getField1());
        if (compareField1 != 0) {
            return -compareField1;
        }
        // For equal field1, sort by field2 ascending
        int compareField2 = item1.getField2().compareTo(item2.getField2());
        if (compareField2 != 0) {
            return compareField2;
        }
        // For equal field1 and field1, sort by item id
        return item1.getId().compareTo(item2.getId());
    }
}

technorati tags:TreeSet, Comparator, Problem