Implement Wavelet Tree with rank and kthSmallest methods

space0032 · space0032 · commit 3ea09f2fbbe2 · 2026-05-06T15:14:06.000+05:30
diff --git a/src/main/java/com/thealgorithms/datastructures/trees/WaveletTree.java b/src/main/java/com/thealgorithms/datastructures/trees/WaveletTree.java
@@ -0,0 +1,243 @@
+package com.thealgorithms.datastructures.trees;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * A Wavelet Tree is a highly efficient data structure used to store sequences
+ * and answer queries like rank, select, and quantile in O(log(max_val - min_val)) time.
+ * This structure is particularly useful in competitive programming and text compression.
+ */
+public class WaveletTree {
+
+    private class Node {
+        int low, high;
+        Node left, right;
+        List<Integer> leftCount; // Prefix sums of elements going to the left child
+
+        /**
+         * Recursively constructs the tree nodes by partitioning the array.
+         *
+         * @param arr  the subarray for the current node
+         * @param low  the minimum possible value in the current node
+         * @param high the maximum possible value in the current node
+         */
+        Node(int[] arr, int low, int high) {
+            this.low = low;
+            this.high = high;
+
+            if (arr.length == 0 || low == high) {
+                return;
+            }
+
+            int mid = low + (high - low) / 2;
+            leftCount = new ArrayList<>(arr.length + 1);
+            leftCount.add(0);
+
+            List<Integer> leftArr = new ArrayList<>();
+            List<Integer> rightArr = new ArrayList<>();
+
+            for (int x : arr) {
+                if (x <= mid) {
+                    leftArr.add(x);
+                    leftCount.add(leftCount.get(leftCount.size() - 1) + 1);
+                } else {
+                    rightArr.add(x);
+                    leftCount.add(leftCount.get(leftCount.size() - 1));
+                }
+            }
+
+            if (!leftArr.isEmpty()) {
+                this.left = new Node(leftArr.stream().mapToInt(i -> i).toArray(), low, mid);
+            }
+            if (!rightArr.isEmpty()) {
+                this.right = new Node(rightArr.stream().mapToInt(i -> i).toArray(), mid + 1, high);
+            }
+        }
+    }
+
+    private Node root;
+    private final int n;
+
+    /**
+     * Constructs a Wavelet Tree from the given array.
+     * The min and max values are determined dynamically from the array.
+     *
+     * @param arr the input array
+     */
+    public WaveletTree(int[] arr) {
+        if (arr == null || arr.length == 0) {
+            this.n = 0;
+            return;
+        }
+        this.n = arr.length;
+        int min = arr[0];
+        int max = arr[0];
+        for (int x : arr) {
+            if (x < min) {
+                min = x;
+            }
+            if (x > max) {
+                max = x;
+            }
+        }
+        root = new Node(arr, min, max);
+    }
+
+    /**
+     * Constructs a Wavelet Tree from the given array with specific min and max values.
+     *
+     * @param arr      the input array
+     * @param minValue the minimum possible value
+     * @param maxValue the maximum possible value
+     */
+    public WaveletTree(int[] arr, int minValue, int maxValue) {
+        if (arr == null || arr.length == 0) {
+            this.n = 0;
+            return;
+        }
+        this.n = arr.length;
+        root = new Node(arr, minValue, maxValue);
+    }
+
+    /**
+     * How many times does the number x appear in the array from index 0 to i (inclusive)?
+     *
+     * @param x the number to search for
+     * @param i the end index (0-based, inclusive)
+     * @return the number of occurrences of x in arr[0...i]
+     */
+    public int rank(int x, int i) {
+        if (root == null || x < root.low || x > root.high || i < 0) {
+            return 0;
+        }
+        // If i is out of bounds, cap it at n - 1
+        int endIdx = Math.min(i, n - 1);
+        return rank(root, x, endIdx + 1);
+    }
+
+    private int rank(Node node, int x, int count) {
+        if (node == null || count == 0) {
+            return 0;
+        }
+        if (node.low == node.high) {
+            return count;
+        }
+        int mid = node.low + (node.high - node.low) / 2;
+        int leftC = node.leftCount.get(count);
+        if (x <= mid) {
+            return rank(node.left, x, leftC);
+        } else {
+            return rank(node.right, x, count - leftC);
+        }
+    }
+
+    /**
+     * What is the 0-based index of the k-th occurrence of the number x in the array?
+     *
+     * @param x the number to search for
+     * @param k the occurrence count (1-based)
+     * @return the 0-based index in the original array, or -1 if x occurs less than k times
+     */
+    public int select(int x, int k) {
+        if (root == null || x < root.low || x > root.high || k <= 0) {
+            return -1;
+        }
+        if (rank(x, n - 1) < k) {
+            return -1;
+        }
+        return select(root, x, 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);
+        }
+    }
+
+    private int binarySearchLeft(List<Integer> prefixSums, int k) {
+        int l = 1, r = prefixSums.size() - 1;
+        int ans = -1;
+        while (l <= r) {
+            int mid = l + (r - l) / 2;
+            if (prefixSums.get(mid) >= k) {
+                ans = mid;
+                r = mid - 1;
+            } else {
+                l = mid + 1;
+            }
+        }
+        return ans == -1 ? -1 : ans - 1; // Convert to 0-based index
+    }
+
+    private int binarySearchRight(List<Integer> prefixSums, int k) {
+        int l = 1, r = prefixSums.size() - 1;
+        int ans = -1;
+        while (l <= r) {
+            int mid = l + (r - l) / 2;
+            if (mid - prefixSums.get(mid) >= k) {
+                ans = mid;
+                r = mid - 1;
+            } else {
+                l = mid + 1;
+            }
+        }
+        return ans == -1 ? -1 : ans - 1; // Convert to 0-based index
+    }
+
+    /**
+     * If you sort the subarray from index left to right, what would be the k-th smallest element?
+     * This query is also commonly known as the quantile query.
+     *
+     * @param left  the start index of the subarray (0-based, inclusive)
+     * @param right the end index of the subarray (0-based, inclusive)
+     * @param k     the rank of the smallest element (1-based, e.g., k=1 is the minimum)
+     * @return the k-th smallest element in the subarray, or -1 if invalid parameters
+     */
+    public int kthSmallest(int left, int right, int k) {
+        if (root == null || left > right || left < 0 || k < 1 || k > right - left + 1) {
+            return -1;
+        }
+        return kthSmallest(root, left, right, 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;
+        }
+
+        int countLeftInLMinus1 = (left == 0) ? 0 : node.leftCount.get(left);
+        int countLeftInR = node.leftCount.get(right + 1);
+        int elementsToLeft = countLeftInR - countLeftInLMinus1;
+
+        if (k <= elementsToLeft) {
+            int newL = countLeftInLMinus1;
+            int newR = countLeftInR - 1;
+            return kthSmallest(node.left, newL, newR, k);
+        } else {
+            int newL = left - countLeftInLMinus1;
+            int newR = right - countLeftInR;
+            return kthSmallest(node.right, newL, newR, k - elementsToLeft);
+        }
+    }
+}
diff --git a/src/test/java/com/thealgorithms/datastructures/trees/WaveletTreeTest.java b/src/test/java/com/thealgorithms/datastructures/trees/WaveletTreeTest.java
@@ -0,0 +1,102 @@
+package com.thealgorithms.datastructures.trees;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import org.junit.jupiter.api.Test;
+
+public class WaveletTreeTest {
+
+    @Test
+    public void testRank() {
+        int[] arr = {5, 1, 2, 5, 1};
+        WaveletTree wt = new WaveletTree(arr);
+
+        // x = 1
+        assertEquals(1, wt.rank(1, 1)); // In [5, 1], '1' appears 1 time
+        assertEquals(2, wt.rank(1, 4)); // In [5, 1, 2, 5, 1], '1' appears 2 times
+        assertEquals(0, wt.rank(1, 0)); // In [5], '1' appears 0 times
+
+        // x = 5
+        assertEquals(1, wt.rank(5, 0)); // In [5], '5' appears 1 time
+        assertEquals(1, wt.rank(5, 2)); // In [5, 1, 2], '5' appears 1 time
+        assertEquals(2, wt.rank(5, 4)); // In [5, 1, 2, 5, 1], '5' appears 2 times
+
+        // Out of bounds / invalid value
+        assertEquals(0, wt.rank(10, 4)); // '10' is not in the array
+        assertEquals(0, wt.rank(5, -1)); // Invalid end index
+    }
+
+    @Test
+    public void testSelect() {
+        int[] arr = {5, 1, 2, 5, 1};
+        WaveletTree wt = new WaveletTree(arr);
+
+        assertEquals(1, wt.select(1, 1)); // 1st '1' is at index 1
+        assertEquals(4, wt.select(1, 2)); // 2nd '1' is at index 4
+
+        assertEquals(0, wt.select(5, 1)); // 1st '5' is at index 0
+        assertEquals(3, wt.select(5, 2)); // 2nd '5' is at index 3
+
+        assertEquals(2, wt.select(2, 1)); // 1st '2' is at index 2
+
+        assertEquals(-1, wt.select(5, 3)); // 3rd '5' doesn't exist
+        assertEquals(-1, wt.select(10, 1)); // '10' doesn't exist
+        assertEquals(-1, wt.select(5, 0)); // invalid k
+    }
+
+    @Test
+    public void testKthSmallest() {
+        int[] arr = {5, 1, 2, 5, 1};
+        WaveletTree wt = new WaveletTree(arr);
+
+        // Array: [5, 1, 2, 5, 1] -> Sorted: [1, 1, 2, 5, 5]
+        assertEquals(1, wt.kthSmallest(0, 4, 1)); // 1st smallest in [5, 1, 2, 5, 1] is 1
+        assertEquals(1, wt.kthSmallest(0, 4, 2)); // 2nd smallest in [5, 1, 2, 5, 1] is 1
+        assertEquals(2, wt.kthSmallest(0, 4, 3)); // 3rd smallest in [5, 1, 2, 5, 1] is 2
+        assertEquals(5, wt.kthSmallest(0, 4, 4)); // 4th smallest in [5, 1, 2, 5, 1] is 5
+        assertEquals(5, wt.kthSmallest(0, 4, 5)); // 5th smallest in [5, 1, 2, 5, 1] is 5
+
+        // Subarray: arr[1..3] = [1, 2, 5] -> Sorted: [1, 2, 5]
+        assertEquals(1, wt.kthSmallest(1, 3, 1)); // 1st smallest in [1, 2, 5] is 1
+        assertEquals(2, wt.kthSmallest(1, 3, 2)); // 2nd smallest in [1, 2, 5] is 2
+        assertEquals(5, wt.kthSmallest(1, 3, 3)); // 3rd smallest in [1, 2, 5] is 5
+
+        // Invalid ranges / arguments
+        assertEquals(-1, wt.kthSmallest(4, 2, 1)); // Invalid range (left > right)
+        assertEquals(-1, wt.kthSmallest(0, 4, 10)); // k > range length
+        assertEquals(-1, wt.kthSmallest(0, 4, 0)); // k < 1
+    }
+
+    @Test
+    public void testEmptyAndSingleElementArray() {
+        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});
+        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 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));
+
+        // Sorted: [-5, -5, -2, 0, 10]
+        assertEquals(-5, wt.kthSmallest(0, 4, 1));
+        assertEquals(-2, wt.kthSmallest(0, 4, 3));
+        assertEquals(10, wt.kthSmallest(0, 4, 5));
+    }
+}
