Java TreeSet 26 ms solution optimized with a custom comparator


  • 3

    The problem with TreeSet solutions is not only that they are O(n log k) when there are O(n) solutions, but also that they search multiple times for the same value. Consider a typical floor / ceiling solution: first it checks for floor and ceiling, then it inserts the new value, and each of those operations is basically a search for the same value! How can we deal with that?

    One idea is to perform the search once, finding the floor, ceiling and insertion point at the same time. Moreover, if the floor or the ceiling are too close, we don't even want to continue the search: we can return true immediately. But how do we do this?

    One way is to create a custom BST implementation. This can get ugly especially if you do RBT rebalancing. But there is a smarter way. Note that TreeSet doesn't insert a value if there is already such value in the set (it is a set, after all). What if we want to achieve the same effect when the value is not exactly equal to some other, but almost equal (within the margin of t)? Why not trick it into thinking that these values are indeed equal? Especially since the docs openly state that TreeSet / TreeMap only relies on compare / Comparator and not on equals when checking for equality.

    Note that our "equality" violates one property of equality: it's not transitive. No big deal, though, because the set will only contain distinct (according to our "equality" criterion) values, and for "less" / "greater" relationships transitivity is preserved.

    So here you go. Still O(n log k) but twice faster than the floor / ceiling one.

    public boolean containsNearbyAlmostDuplicate(int[] nums, int k, final int t) {
        if (k == 0 || nums.length <= 1) {
            return false;
        }
        if (k >= nums.length) {
            k = nums.length - 1; // note: mutating a formal parameter!
        }
        NavigableSet<Integer> kSet = new TreeSet<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer i1, Integer i2) {
                long diff = (long) i1 - i2; // long because can overflow
                if (Math.abs(diff) <= t) {
                    return 0;
                } else {
                    return diff > 0 ? +1 : -1;
                }
            }
        });
        for (int i = 0; i <= k; ++i) {
            if (!kSet.add(nums[i])) {
                return true;
            }
        }
        for (int i = k + 1; i < nums.length; ++i) {
            kSet.remove(nums[i - k - 1]);
            if (!kSet.add(nums[i])) {
                return true;
            }
        }
        return false;
    }

  • 0
    X
    public boolean containsNearbyAlmostDuplicate(int[] nums, int k, int t) {
        if (k == 0 || nums.length <= 1) return false;
        TreeSet<Integer> kSet = new TreeSet<>((i1, i2)->{
                long diff = (long) i1 - i2;
                if (Math.abs(diff) <= t) return 0;
                else return diff > 0 ? +1 : -1;
            });
    
        for (int i = 0; i < nums.length; ++i) {
            if (!kSet.add(nums[i])) return true;
            if(i>=k)kSet.remove(nums[i - k]);
        }
        return false;      
    }

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.