Python Program for Wavelet Trees

Wavelet Trees are advanced data structures used to efficiently answer range queries on arrays or strings.

They are especially useful for queries like finding the kth smallest element, counting occurrences, and rank queries.

1. Understanding the Problem

Given an array, efficiently answer queries like kth smallest element in a range or frequency of an element.

Input: arr = [3, 1, 4, 1, 5, 9, 2]
Query: kth smallest in range [1,5], k=3
Output: 3

2. Method 1: Basic Concept

              
            
Python

          Conceptual structure
        
# Wavelet Tree splits values into ranges
# Left child: smaller values
# Right child: larger values

# Example structure representation
# Root: [3,1,4,1,5,9,2]
# Left: [3,1,1,2]
# Right: [4,5,9]

The array is recursively divided based on value ranges.

3. Method 2: Build Wavelet Tree

              
            
Python

          Wavelet tree construction
        
class WaveletTree:
    def __init__(self, data, low, high):
        self.low = low
        self.high = high
        self.left = None
        self.right = None
        self.map = []

        if low == high or not data:
            return

        mid = (low + high) // 2
        left_data = []
        right_data = []

        for num in data:
            if num <= mid:
                left_data.append(num)
                self.map.append(1)
            else:
                right_data.append(num)
                self.map.append(0)

        self.left = WaveletTree(left_data, low, mid)
        self.right = WaveletTree(right_data, mid+1, high)

This builds the tree by dividing elements into ranges.

4. Method 3: Kth Smallest Query

              
            
Python

          Find kth smallest
        
def kth(self, l, r, k):
    if self.low == self.high:
        return self.low

    in_left = sum(self.map[l:r+1])

    if k <= in_left:
        new_l = sum(self.map[:l])
        new_r = sum(self.map[:r+1]) - 1
        return self.left.kth(new_l, new_r, k)
    else:
        new_l = l - sum(self.map[:l])
        new_r = r - sum(self.map[:r+1]) + 1
        return self.right.kth(new_l, new_r, k - in_left)

This efficiently finds kth smallest element in a range.

5. Method 4: Rank Query

              
            
Python

          Count occurrences
        
def rank(self, l, r, x):
    if l > r or x < self.low or x > self.high:
        return 0
    if self.low == self.high:
        return r - l + 1

    mid = (self.low + self.high) // 2

    left_l = sum(self.map[:l])
    left_r = sum(self.map[:r+1]) - 1

    if x <= mid:
        return self.left.rank(left_l, left_r, x)
    else:
        right_l = l - left_l
        right_r = r - left_r - 1
        return self.right.rank(right_l, right_r, x)

Rank query counts frequency of an element in a range.

6. Method 5: Handling Edge Cases

              
            
Python

          Edge cases
        
arr = list(map(int, input().split()))

if not arr:
    print("Empty array")
else:
    tree = WaveletTree(arr, min(arr), max(arr))
    print("Tree built successfully")

Handles empty arrays and initialization safely.

7. Algorithm

1. Divide array based on value range.

2. Store mapping of elements.

3. Recursively build left and right subtrees.

4. Use mapping to answer queries.

5. Support kth, rank, and range queries.

8. Common Mistakes

1. Incorrect mapping array handling.

2. Wrong index calculations.

3. Not handling empty inputs.

4. Confusing value range with index range.

9. Applications

1. Range queries (kth smallest).

2. Competitive programming.

3. Database indexing.

4. Data analytics.

Conclusion

Wavelet Trees are powerful for solving complex range query problems efficiently.

They are widely used in competitive programming and advanced data processing systems.