Implement Wavelet Tree with rank and kthSmallest methods

Author: space0032

src/main/java/com/thealgorithms/datastructures/trees/WaveletTree.java

@@ -26,7 +26,7 @@ private class Node {
             this.low = low;
             this.high = high;
 
-            if (arr.length == 0 || low == high) {
+            if (low == high) {
                 return;
             }
 
@@ -150,24 +150,15 @@ public int select(int x, int k) {
     }
 
     private int select(Node node, int x, int k) {
-        if (node == null) {
-            return -1;
-        }
         if (node.low == node.high) {
             return k - 1; // 0-based index within the imaginary array at the leaf
         }
         int mid = node.low + (node.high - node.low) / 2;
         if (x <= mid) {
             int posInLeft = select(node.left, x, k);
-            if (posInLeft == -1) {
-                return -1;
-            }
             return binarySearchLeft(node.leftCount, posInLeft + 1);
         } else {
             int posInRight = select(node.right, x, k);
-            if (posInRight == -1) {
-                return -1;
-            }
             return binarySearchRight(node.leftCount, posInRight + 1);
         }
     }
@@ -219,9 +210,6 @@ public int kthSmallest(int left, int right, int k) {
     }
 
     private int kthSmallest(Node node, int left, int right, int k) {
-        if (node == null) {
-            return -1;
-        }
         if (node.low == node.high) {
             return node.low;
         }

src/test/java/com/thealgorithms/datastructures/trees/WaveletTreeTest.java

@@ -69,27 +69,42 @@ public void testKthSmallest() {
 
     @Test
     public void testEmptyAndSingleElementArray() {
-        WaveletTree wtEmpty = new WaveletTree(new int[]{});
+        WaveletTree wtEmpty = new WaveletTree(new int[] {});
         assertEquals(0, wtEmpty.rank(1, 0));
         assertEquals(-1, wtEmpty.select(1, 1));
         assertEquals(-1, wtEmpty.kthSmallest(0, 0, 1));
 
-        WaveletTree wtSingle = new WaveletTree(new int[]{42});
+        WaveletTree wtSingle = new WaveletTree(new int[] {42});
         assertEquals(1, wtSingle.rank(42, 0));
         assertEquals(0, wtSingle.rank(42, -1));
         assertEquals(0, wtSingle.select(42, 1));
         assertEquals(-1, wtSingle.select(42, 2));
         assertEquals(42, wtSingle.kthSmallest(0, 0, 1));
     }
 
+    @Test
+    public void testNullArrayAndCustomBounds() {
+        WaveletTree wtNull = new WaveletTree(null);
+        assertEquals(0, wtNull.rank(1, 0));
+
+        WaveletTree wtNullCustom = new WaveletTree(null, 1, 5);
+        assertEquals(-1, wtNullCustom.select(1, 1));
+
+        int[] arr = {5, 1, 2, 5, 1};
+        WaveletTree wtCustom = new WaveletTree(arr, 1, 10);
+        assertEquals(2, wtCustom.rank(5, 4));
+        assertEquals(0, wtCustom.rank(4, 4)); // Query an element inside bounds but not in array
+        assertEquals(0, wtCustom.rank(10, 4)); // Query upper bound
+    }
+
     @Test
     public void testNegativeValues() {
         int[] arr = {-5, 10, -2, 0, -5};
         WaveletTree wt = new WaveletTree(arr);
 
         assertEquals(2, wt.rank(-5, 4));
         assertEquals(1, wt.rank(0, 3));
-        
+
         assertEquals(0, wt.select(-5, 1));
         assertEquals(4, wt.select(-5, 2));
         assertEquals(3, wt.select(0, 1));