Ordering is not correct in TreeSet

Question

I'm trying to solve https://leetcode.com/problems/lru-cache/description/.

I'm using treeset to save unique keys and order them based on when it's inserted. The treeset is not returning elements in the way I specified using Camparator.

class LRUCache {

    Map<Integer, Integer> keyValue;// ()
    TreeSet<Pair> keys; //key, timestamp (1,-1000)
    int capacity;//2
    int timestamp;
    public LRUCache(int capacity) {
        keyValue = new HashMap<>(capacity);
        keys = new TreeSet<>((Pair a,Pair b) -> {
            System.out.println("inside comparator implementation");
            int result = a.x == b.x ? 0 : a.y - b.y == 0 ? -1 : a.y - b.y;
            System.out.println("result: " + result);
            return result;
        });
        this.capacity = capacity;
    }

    public int get(int key) {

        timestamp++;
        Pair pair = new Pair(key, timestamp);
        boolean found = keys.contains(pair);
        if (found) {
            System.out.println("keys contains " + key + ". So, removing it to update");
            keys.remove(pair);
        } else {
            System.out.println("keys doesn't contains " + key + ". So, not removing it. Directly adding it");
        }

        boolean added = keys.add(pair);
        System.out.println("add status: " + added + " for: " + pair);

        return keyValue.getOrDefault(key, -1);
    }

    public void put(int key, int value) {
        int priority = Integer.MIN_VALUE + timestamp++;
        Pair pair = new Pair(key, priority);
        System.out.println("key: " + key + " priority: " + priority);
        boolean found = keys.contains(pair);
        if (!found) {
            System.out.println("keys doesn't " + key + " and " + value + ". So, adding to keys");
            keys.add(pair);
        } else {
            System.out.println("keys contains " + key + " and " + value + ". So, not adding to keys");
        }

        keyValue.put(key, value);
        if (keys.size() > capacity) {
            System.out.println("size > capacity");
            Pair remove = keys.first();
            System.out.println("removing: " + remove);
            keys.remove(remove);
            keyValue.remove(remove.x);
        }
    }

    public static void main(String[] args) {
        LRUCache lruCache = new LRUCache(2);
        System.out.println();
        System.out.println("putting 1,1");
        lruCache.put(1, 1);

        System.out.println();
        System.out.println("putting 2,2");
        lruCache.put(2, 2);

        System.out.println();
        System.out.println("getting 1");
        System.out.println(lruCache.get(1));

        System.out.println();
        System.out.println("c");
        lruCache.put(3, 3);

        System.out.println();
        System.out.println("getting 2");
        System.out.println(lruCache.get(2));
    }
}

class Pair{
    int x;
    int y;
    public Pair(int x, int y) {
        this.x = x;
        this.y = y;
    }

    public boolean equals(Object o) {
        System.out.println("inside pair equals");
        if (this == o) {
            return true;
        }
        if (!(o instanceof Pair)) {
            return false;
        }
        return this.x == ((Pair) o).x;
    }

    public int hashCode() {
        return x;
    }

    @Override
    public String toString() {
        return "Pair{" +
                "x=" + x +
                ", y=" + y +
                '}';
    }
}

I'm expecting (2, -2147483647) to be dropped after adding (3,3), since (2, -2147483647) has negative timestamp.

This is the set that is printed before I expected (2, -2147483647) is removed. [Pair{x=1, y=3}, Pair{x=2, y=-2147483647}, Pair{x=3, y=-2147483645}]

Why is Pair{x=1, y=3} first compared to {x=2, y=-2147483647}. I wrote the comparator to order in treeset based on y values, why is it working as expected?

I tried creating a TreeSetPractise

public class TreeSetPractise {
    public static void main(String[] args) {
        TreeSet<Pair> ts = new TreeSet<>(( a,  b) -> {
            System.out.println("inside comparator implementation");
            int result = a.x == b.x ? 0 : a.y - b.y == 0 ? -1 : a.y - b.y;
            System.out.println("result: " + result);
            return result;
        });
        ts.add(new Pair(1,10));
        ts.add(new Pair(2, -20));
        ts.add(new Pair(3, 30));
        System.out.println(ts);
        for (Pair i : ts) {
            System.out.println(i);
        }
    }
}

class Pair {
    int x;
    int y;
    public Pair(int x, int y) {
        this.x = x;
        this.y = y;
    }

    @Override
    public String toString() {
        return "Pair{" +
                "x=" + x +
                ", y=" + y +
                '}';
    }
}

This is working as expected.

The main problem here is, for the elements of TreeSet:Pair{x,y} has to check equality on x and ordering on y. Meaning Set functionality is based on x and Tree functionality is based on y.
Update: Most of the comments/answers points out here that my comparator is wrong. The only problem with it might be overflow, but there is no problem with logic. Also, the bigger problem for me is, TreeSet is considering two objects to be equal if comparator is returning 0. Shouldn't it call equals method on if compareTo returns 0?

Answer 1

The Answer by rph correctly diagnoses your problem. I'll add a few more points.

Harmonize your equals/hashCode with your Comparator

As Comments discussed, the logic of your override of equals (and hashCode) must be in harmony with your Comparator.

Record

You could more briefly declare your Pair class as a record.

In a record, by default, the compiler implicitly creates an overriding equals and hashCode implementation that examines each and every member field. So less work for you.

record Point ( int x , int y ) {}

Comparator builder

As commented, avoid your errors by building a Comparator via the convenience methods (…comparing…). See comparingInt.

Important: Notice in code below that our Comparator implementation considers each and every field (x & y), to be consistent with the default equals (and hashCode) implementations of our record.

By the way, you can refer to your TreeSet as the more general NavigableSet.

NavigableSet < Point > ts = new TreeSet <>(
        Comparator
                .comparingInt( Point :: x )
                .thenComparingInt( Point :: y )
);

As of Java 21, SequencedSet is the more general super-interface of NavigableSet/SortedSet. See JEP 431. Even more general is SequencedCollection.

SequencedSet <Pair> ts = new TreeSet<> …

Complete example

package work.basil.example.collections;

import java.util.*;

public class TreeSetOrdering
{
    public static void main ( String[] args )
    {
        record Point( int x , int y ) { }

        SequencedSet < Point > ts = new TreeSet <>(
                Comparator
                        .comparingInt( Point :: x )
                        .thenComparingInt( Point :: y )
        );

        ts.addAll(
                List.of(
                        new Point( 7 , 42 ) ,
                        new Point( 11 , 22 ) ,
                        new Point( 1 , 2 )
                )
        );
        System.out.println( "ts = " + ts );
    }
}

ts = [Point[x=1, y=2], Point[x=7, y=42], Point[x=11, y=22]]

Answer 2

Looking at the comparator logic:

int result = a.x == b.x ? 0 : a.y - b.y == 0 ? -1 : a.y - b.y;

There are a few issues with a.y - b.y == 0 ? -1 : a.y - b.y.

1. if a.y equals b.y then it considers a < b, which is likely not intended.
2. otherwise performs an arithmetic operation, which might overflow, and thus provide the wrong result, when b.y can be a large negative number like Integer.MIN_VALUE.

If the goal is to compare y if x is not the same, you can replace that part with a.y == b.y ? 0 : (a.y < b.y) ? - 1 : 1. However, a safer (less error prone) approach is to just use Integer.compare to compare a.y and b.y.

int result = a.x == b.x ? 0 : Integer.compare(a.y, b.y);

Answer 3

Oh boy. Let's start with the big issues:

1. The problem statement requires get and put be O(1) operations, but TreeSet.put() and TreeSet.get() take O(log n) time. That is, a TreeSet is too slow.
2. If you order your TreeSet by key first, and age second, you can lookup by key in O(log n), but key.first() is the smallest rather than the oldest entry, which causes the wrong key to be evicted. If you order by age first, and key second, you can not lookup by key. There is no odering that supports both efficient lookup by key (regardless of age), and efficient lookup by age (regardless of key). You can have either, but not both. You could use two different TreeMaps - one ordered by age, one ordered by key, that map the age to the key, and the key to the age, but that gets rather complex.

And finally, your code is rather complicated for what it does.

The simple way to implement a LRUCache in Java is this:

class LRUCache<K,V> extends LinkedHashMap<K,V> {
    final int capacity;

    LRUCache(int capacity) {
        super(10, 0.75f, true);
        this.capacity = capacity;
    }
    
    @Override
    protected boolean removeEldestEntry(Entry<K, V> eldest) {
        return size() > capacity;
    }
}

As you can see from the class name of the super class, the data is kept in a map for fast lookup, with a linked list for keeping track of access order. To maintain the list in access order, we remove the node on access (which is O(1)), and then add it to the tail of the list (which is O(1)).