Wavelet Matrix
The Wavelet Matrix is a compact data structure that operates on a sequence of integers . Each integer is typically drawn from an alphabet of size , often represented as integers in the range . The structure is particularly efficient for answering various types of queries on the sequence, such as rank, select, and range queries.

You are encouraged to solve this task according to the task description, using any language you may know.
- Wavelet Matrix Problem
Your task is to implement a data structure (or use a conceptual Wavelet Matrix) that, after processing an initial sequence , can efficiently answer a series of queries.
- Sequence and Construction Context
The Wavelet Matrix is conceptually built by considering the binary representations of the numbers in the sequence . It consists of levels, one for each bit position from the most significant bit (MSB) to the least significant bit (LSB). At each level (from to , corresponding to the -th bit from MSB):
- A bitvector is constructed. For each number in the current (possibly permuted) sequence, if its -th bit is 1, a 1 is appended to ; otherwise, a 0 is appended.
- The numbers are then stably reordered for the next level: all numbers for which the -th bit was 0 appear first (in their original relative order), followed by all numbers for which the bit was 1 (also in their original relative order).
Each bitvector is typically augmented to support fast rank operations (counting 0s or 1s up to a certain position).
...which is:
As with the FreeBASIC sample, this sample is split into two parts - an "include" file containing the main Wavelet Matrix code and a separate "main" routine.
Additionally, the following libraries are used:
The source of these libraries are on separate pages on Rosetta Code - see the above links.
See the above links for details on how "include" files can be used with your program. particularly if you aren't using ALGOL 68 Genie.
File: wavelet_matrix.incl.a68:
# wavelet_matrix.incl.a68: Translated from the FreeBASIC sample #
# BitRank class for rank operations #
MODE BITRANK = STRUCT( REF[]BITS bitr # Array to store bits #
, REF[]INT count # Array to store counts #
);
# WaveletMatrix class #
MODE WAVELETMATRIX = STRUCT( REF[]BITRANK b # Array of BitRanks #
, REF[]INT posic # posicition array #
, INT height # Height of the matrix #
);
# Get bit at position bn from value, NB in Algol 68, the leftmost bit is #
# bit 1, the rightmost is numbered bits width #
PROC get bit at = ( INT value, bn )INT: ABS ( ( bits width - bn ) ELEM BIN value );
# Count number of 1's in [0, i) #
PROC rank1 = ( REF BITRANK br, INT i )INT:
BEGIN
INT block index := i SHR 6;
INT bitposic := i AND 63;
INT count := 0;
# Add count from previous blocks #
IF block index > 0 THEN count := ( count OF br )[ block index - 1 ] FI;
# Count bits in current block #
FOR j FROM 0 TO bitposic - 1 DO
IF ( ( bitr OF br )[ block index ] AND ( 1 SHL j ) ) /= 16r0 THEN count +:= 1 FI
OD;
count
END # rank1 # ;
# Count number of 0's in [0, i) #
PROC rank0 = ( REF BITRANK br, INT i )INT: i - rank1( br, i );
# Initialize a BitRank structure #
PROC init bit rank = ( REF BITRANK br, INT size )VOID:
BEGIN
INT num blocks = ( size + 63 ) OVER 64;
HEAP[ 0 : num blocks ]BITS new bits;
HEAP[ 0 : num blocks ]INT new counts;
bitr OF br := new bits;
count OF br := new counts;
FOR i FROM 0 TO num blocks DO
( bitr OF br )[ i ] := 16r0;
( count OF br )[ i ] := 0
OD
END # init bit rank # ;
# Initialize the WaveletMatrix with a vector #
PROC create wavelet matrix = ( []INT vec, INT n, max val )WAVELETMATRIX:
BEGIN
# Calculate height based on maximum value #
WAVELETMATRIX wm;
height OF wm := 0;
INT tmp := max val;
WHILE tmp > 0 DO
height OF wm +:= 1;
tmp OVERAB 2
OD;
# Initialize arrays #
HEAP[ 0 : height OF wm - 1 ]BITRANK new b;
HEAP[ 0 : height OF wm - 1 ]INT new posic;
b OF wm := new b;
posic OF wm := new posic;
# Create a copy of the input vector for manipulation #
[ 0 : n - 1 ]INT tmp vec := vec[ AT 0 ];
FOR h FROM 0 TO height OF wm - 1 DO
# Initialize BitRank #
init bit rank( ( b OF wm )[ h ], n );
# Count zeros #
INT zeros := 0;
FOR i FROM 0 TO n - 1 DO
IF get bit at( tmp vec[ i ], height OF wm - h - 1 ) = 0 THEN
zeros +:= 1
ELSE
# Set bit in BitRank #
INT block idx = i SHR 6;
INT bitposic = i AND 63;
( bitr OF ( b OF wm )[ h ] )[ block idx ] ORAB ( 1 SHL bitposic )
FI
OD;
# Build count array #
INT running count := 0;
FOR i FROM 0 TO UPB bitr OF ( b OF wm )[ h ] DO
INT bit count := 0;
BITS bitr := ( bitr OF ( b OF wm )[ h ] )[ i ];
# Count bits in this block #
WHILE bitr /= 16r0 DO
IF ( bitr AND 1 ) /= 16r0 THEN bit count +:= 1 FI;
bitr SHRAB 1
OD;
( count OF ( b OF wm )[ h ] )[ i ] := running count;
running count +:= bit count
OD;
# Stable partition - separate 0's and 1's #
[ 0 : zeros - 1 ]INT zero values;
[ 0 : n - zeros - 1 ]INT one values;
INT zero idx := 0, one idx := 0;
FOR i FROM 0 TO n - 1 DO
IF get bit at( tmp vec[ i ], height OF wm - h - 1 ) = 0 THEN
zero values[ zero idx ] := tmp vec[ i ];
zero idx +:= 1
ELSE
one values[ one idx ] := tmp vec[ i ];
one idx +:= 1
FI
OD;
# Combine arrays back #
tmp vec[ 0 : zeros - 1 ] := zero values[ 0 : zeros - 1 ];
tmp vec[ zeros : n - 1 ] := one values[ 0 : n - zeros - 1 ];
( posic OF wm )[ h ] := zeros
OD;
wm
END # create wavelet matrix # ;
# kth smallest element in [l, r) #
PROC quantile = ( REF WAVELETMATRIX wm, INT k in, l in, r in )INT:
BEGIN
INT k := k in, l := l in, r := r in;
INT res := 0;
FOR i FROM 0 TO height OF wm - 1 DO
INT zeros = rank0( ( b OF wm )[ i ], r )
- rank0( ( b OF wm )[ i ], l )
;
IF zeros > k THEN # Go to left (0) child #
l := rank0( ( b OF wm )[ i ], l );
r := rank0( ( b OF wm )[ i ], r )
ELSE # Go to right (1) child #
l := ( posic OF wm )[ i ] + rank1( ( b OF wm )[ i ], l );
r := ( posic OF wm )[ i ] + rank1( ( b OF wm )[ i ], r );
k -:= zeros;
res ORAB ( 1 SHL ( height OF wm - i - 1 ) )
FI
OD;
res
END # quantile # ;
# end wavelet_matrix.incl.a68 #Main routine
BEGIN # Wavelet Matrix - translated from the FreeBASIC sample #
PR read "bits.incl.a68" PR # include bit-manipulation utilities #
PR read "rows.incl.a68" PR # include row (array) related utilities #
PR read "sort.incl.a68" PR # include sort utilities #
PR read "wavelet_matrix.incl.a68" PR # WAVELETMATRIX and related code #
BEGIN # task #
# Define the array and queries #
[]INT a = ( 3374, 956, 2114, 3415, 3437 );
INT n = UPB a;
# Create a sorted copy with unique values #
[ 0 : n - 1 ]INT sorted a := a[ AT 0 ];
QUICKSORT sorted a;
INT unique count := 1; # Count unique values #
FOR i FROM LWB sorted a + 1 TO UPB sorted a DO
IF sorted a[ i ] /= sorted a[ i - 1 ] THEN unique count +:= 1 FI
OD;
# Create array with unique values #
[ 0 : unique count - 1 ]INT unique a;
unique a[ 0 ] := sorted a[ 0 ];
INT idx := 0;
FOR i TO n - 1 DO
IF sorted a[ i ] /= sorted a[ i - 1 ] THEN
unique a[ idx +:= 1 ] := sorted a[ i ]
FI
OD;
# Map original values to indices in unique array #
[ 0 : n - 1 ]INT mapped;
FOR i FROM LWB mapped TO UPB mapped DO
mapped[ i ] := unique a INDEXOF a[ AT 0 ][ i ]
OD;
# Create wavelet matrix #
WAVELETMATRIX wm := create wavelet matrix( mapped, n, unique count );
# Process queries #
[,]INT queries = ( ( 2, 2, 1 ), ( 3, 4, 1 ), ( 4, 5, 1 ), ( 1, 2, 2 ), ( 4, 4, 1 ) );
# ^^^^^^^ bounds are [ 1 : 5, 1 : 3 ] #
FOR i FROM LWB queries TO UPB queries DO
INT l = queries[ i, 1 ] - 1; # Convert to 0-indexed #
INT r = queries[ i, 2 ];
INT k = queries[ i, 3 ];
print( ( " ", whole( unique a[ quantile( wm, k - 1, l, r ) ], 0 ) ) )
OD;
print( ( newline ) )
END
END- Output:
956 2114 3415 3374 3415
using System;
using System.Collections.Generic;
using System.Linq;
public class WaveletMatrixDemo
{
// BitRank is a rank data structure for bit vectors
private class BitRank
{
private long[] block;
private int[] count;
// Resize resizes the bit vector to the given length
public void Resize(int num)
{
block = new long[((num + 1) >> 6) + 1];
count = new int[block.Length];
}
// Set sets bit at position i
public void Set(int i, int val)
{
if (val == 1)
{
block[i >> 6] |= (1L << (i & 63));
}
}
// Build builds the rank structure
public void Build()
{
for (int i = 1; i < block.Length; i++)
{
count[i] = count[i - 1] + PopcountLL(block[i - 1]);
}
}
// PopcountLL counts number of 1's in a 64-bit integer
private int PopcountLL(long n)
{
return BitOperations.PopCount(n);
}
// Rank1 counts number of 1's in [0, i)
public int Rank1(int i)
{
return count[i >> 6] + PopcountLL(block[i >> 6] & ((1L << (i & 63)) - 1));
}
// Rank1FromTo counts number of 1's in [i, j)
public int Rank1FromTo(int i, int j)
{
return Rank1(j) - Rank1(i);
}
// Rank0 counts number of 0's in [0, i)
public int Rank0(int i)
{
return i - Rank1(i);
}
// Rank0FromTo counts number of 0's in [i, j)
public int Rank0FromTo(int i, int j)
{
return Rank0(j) - Rank0(i);
}
}
// WaveletMatrix is a wavelet matrix data structure
private class WaveletMatrix
{
private int height;
private BitRank[] B;
private int[] pos;
// Constructor creates a new wavelet matrix
public WaveletMatrix(int[] vec, params int[] sigma)
{
int s = 0;
if (sigma.Length > 0)
{
s = sigma[0];
}
else
{
// Find the maximum element and use that as sigma
foreach (int v in vec)
{
if (v > s)
{
s = v;
}
}
s++;
}
Init(vec, s);
}
private void Init(int[] vec, int sigma)
{
// Calculate height based on sigma value
if (sigma == 1)
{
height = 1;
}
else
{
height = 64 - BitOperations.LeadingZeroCount((ulong)(sigma - 1));
}
B = new BitRank[height];
pos = new int[height];
for (int i = 0; i < height; i++)
{
B[i] = new BitRank();
B[i].Resize(vec.Length);
for (int j = 0; j < vec.Length; j++)
{
B[i].Set(j, Get(vec[j], height - i - 1));
}
B[i].Build();
// Use a local variable to capture the current i value
int currentLevel = i;
pos[i] = StablePartition(vec, c => Get(c, height - currentLevel - 1) == 0);
}
}
// StablePartition is equivalent to C++ stable_partition
private int StablePartition(int[] arr, Func<int, bool> predicate)
{
List<int> result = new List<int>(arr.Length);
List<int> falseValues = new List<int>(arr.Length);
foreach (int item in arr)
{
if (predicate(item))
{
result.Add(item);
}
else
{
falseValues.Add(item);
}
}
int partitionPoint = result.Count;
result.AddRange(falseValues);
// Update the original array
for (int i = 0; i < result.Count; i++)
{
arr[i] = result[i];
}
return partitionPoint;
}
// Get returns bit at position i from val
private int Get(int val, int i)
{
return (val >> i) & 1;
}
// Rank counts occurrences of val in range [l, r)
public int Rank(int val, int l, int r)
{
return RankSingle(val, r) - RankSingle(val, l);
}
// RankSingle counts occurrences of val in range [0, i)
public int RankSingle(int val, int i)
{
int p = 0;
for (int j = 0; j < height; j++)
{
if (Get(val, height - j - 1) == 1)
{
p = pos[j] + B[j].Rank1(p);
i = pos[j] + B[j].Rank1(i);
}
else
{
p = B[j].Rank0(p);
i = B[j].Rank0(i);
}
}
return i - p;
}
// Quantile returns kth smallest element in [l, r)
public int Quantile(int k, int l, int r)
{
int res = 0;
for (int i = 0; i < height; i++)
{
int j = B[i].Rank0FromTo(l, r);
if (j > k)
{
l = B[i].Rank0(l);
r = B[i].Rank0(r);
}
else
{
l = pos[i] + B[i].Rank1(l);
r = pos[i] + B[i].Rank1(r);
k -= j;
res |= (1 << (height - i - 1));
}
}
return res;
}
// RangeFreq counts elements in [l, r) that are in value range [a, b)
public int RangeFreq(int l, int r, int a, int b)
{
return RangeFreqRecursive(l, r, a, b, 0, 1 << height, 0);
}
private int RangeFreqRecursive(int i, int j, int a, int b, int l, int r, int x)
{
if (i == j || r <= a || b <= l)
{
return 0;
}
int mid = (l + r) >> 1;
if (a <= l && r <= b)
{
return j - i;
}
else
{
int left = RangeFreqRecursive(
B[x].Rank0(i),
B[x].Rank0(j),
a, b, l, mid, x + 1
);
int right = RangeFreqRecursive(
pos[x] + B[x].Rank1(i),
pos[x] + B[x].Rank1(j),
a, b, mid, r, x + 1
);
return left + right;
}
}
// RangeMin finds minimum value in [l, r) within value range [a, b), -1 if not found
public int RangeMin(int l, int r, int a, int b)
{
return RangeMinRecursive(l, r, a, b, 0, 1 << height, 0, 0);
}
private int RangeMinRecursive(int i, int j, int a, int b, int l, int r, int x, int val)
{
if (i == j || r <= a || b <= l)
{
return -1;
}
if (r - l == 1)
{
return val;
}
int mid = (l + r) >> 1;
int res = RangeMinRecursive(
B[x].Rank0(i),
B[x].Rank0(j),
a, b, l, mid, x + 1, val
);
if (res < 0)
{
return RangeMinRecursive(
pos[x] + B[x].Rank1(i),
pos[x] + B[x].Rank1(j),
a, b, mid, r, x + 1,
val + (1 << (height - x - 1))
);
}
else
{
return res;
}
}
}
// Binary search to find index in sorted array
private static int Find(int[] arr, int x)
{
int left = 0;
int right = arr.Length;
while (left < right)
{
int mid = (left + right) / 2;
if (arr[mid] < x)
{
left = mid + 1;
}
else
{
right = mid;
}
}
return left;
}
// Custom BitOperations class to replicate Java's bit counting functionality
private static class BitOperations
{
public static int PopCount(long x)
{
// Implementation of population count (count of set bits)
ulong value = (ulong)x;
int count = 0;
while (value != 0)
{
count += (int)(value & 1);
value >>= 1;
}
return count;
}
public static int LeadingZeroCount(ulong x)
{
// Count the number of leading zeros in a 64-bit integer
if (x == 0) return 64;
int count = 0;
// Test the highest 32 bits
if ((x & 0xFFFFFFFF00000000UL) == 0)
{
count += 32;
x <<= 32;
}
// Test the highest 16 bits
if ((x & 0xFFFF000000000000UL) == 0)
{
count += 16;
x <<= 16;
}
// Test the highest 8 bits
if ((x & 0xFF00000000000000UL) == 0)
{
count += 8;
x <<= 8;
}
// Test the highest 4 bits
if ((x & 0xF000000000000000UL) == 0)
{
count += 4;
x <<= 4;
}
// Test the highest 2 bits
if ((x & 0xC000000000000000UL) == 0)
{
count += 2;
x <<= 2;
}
// Test the highest bit
if ((x & 0x8000000000000000UL) == 0)
{
count += 1;
}
return count;
}
}
public static void Main(string[] args)
{
int n = 5;
int[] a = { 3374, 956, 2114, 3415, 3437 };
int[] input = new int[n];
Array.Copy(a, input, n);
int[] backup = new int[n];
Array.Copy(a, backup, n);
// Sort and deduplicate the array
int[] sortedA = new int[n];
Array.Copy(a, sortedA, n);
Array.Sort(sortedA);
// Deduplicate
List<int> uniqueAList = new List<int>();
for (int i = 0; i < sortedA.Length; i++)
{
if (i == 0 || sortedA[i] != sortedA[i - 1])
{
uniqueAList.Add(sortedA[i]);
}
}
// Convert List to array
int[] uniqueA = uniqueAList.ToArray();
// Map original values to their indices in the unique array
for (int i = 0; i < n; i++)
{
input[i] = Find(uniqueA, backup[i]);
}
int[][] lrkVector = new int[][]
{
new int[] { 2, 2, 1 },
new int[] { 3, 4, 1 },
new int[] { 4, 5, 1 },
new int[] { 1, 2, 2 },
new int[] { 4, 4, 1 }
};
WaveletMatrix wm = new WaveletMatrix(input);
foreach (int[] lrk in lrkVector)
{
int l = lrk[0];
int r = lrk[1];
int k = lrk[2];
l--; // Convert to 0-indexed
Console.WriteLine(uniqueA[wm.Quantile(k - 1, l, r)]);
}
}
}
- Output:
956 2114 3415 3374 3415
REM wavelet_matrix.bas
' BitRank class for rank operations
Type BitRank
bits(Any) As Ulong ' Array to store bits
count(Any) As Integer ' Array to store counts
End Type
' WaveletMatrix class
Type WaveletMatrix
B(Any) As BitRank ' Array of BitRanks
posic(Any) As Integer ' posicition array
height As Integer ' Height of the matrix
End Type
' Get bit at position i from value
Function GetBitAt(value As Integer, i As Integer) As Integer
Return (value Shr i) And 1
End Function
' Count number of 1's in [0, i)
Function Rank1(Byref br As BitRank, i As Integer) As Integer
Dim As Integer blockIndex = i Shr 6
Dim As Integer bitposic = i And 63
Dim As Integer cnt = 0
' Add cnt from previous blocks
If blockIndex > 0 Then cnt = br.count(blockIndex - 1)
' Count bits in current block
For j As Integer = 0 To bitposic - 1
If (br.bits(blockIndex) And (1ULL Shl j)) <> 0 Then cnt += 1
Next
Return cnt
End Function
' Count number of 0's in [0, i)
Function Rank0(Byref br As BitRank, i As Integer) As Integer
Return i - Rank1(br, i)
End Function
' Initialize a BitRank structure
Sub InitBitRank(Byref br As BitRank, size As Integer)
Dim As Integer numBlocks = (size + 63) \ 64
Redim br.bits(0 To numBlocks)
Redim br.count(0 To numBlocks)
' Initialize arrays
For i As Integer = 0 To numBlocks
br.bits(i) = 0
br.count(i) = 0
Next
End Sub
' Initialize the WaveletMatrix with a vector
Sub CreateWaveletMatrix(Byref wm As WaveletMatrix, vec() As Integer, n As Integer, maxVal As Integer)
' Calculate height based on maximum value
wm.height = 0
Dim As Integer i, tmp = maxVal
While tmp > 0
wm.height += 1
tmp Shr= 1
Wend
' Initialize arrays
Redim wm.B(0 To wm.height - 1)
Redim wm.posic(0 To wm.height - 1)
' Create a copy of the input vector for manipulation
Dim As Integer tmpVec(n - 1)
For i = 0 To n - 1
tmpVec(i) = vec(i)
Next
For h As Integer = 0 To wm.height - 1
' Initialize BitRank
InitBitRank(wm.B(h), n)
'Redim wm.B(h).bits(0 To (n + 63) \ 64)
'Redim wm.B(h).cnt(0 To (n + 63) \ 64)
' Count zeros
Dim As Integer zeros = 0
For i = 0 To n - 1
If GetBitAt(tmpVec(i), wm.height - h - 1) = 0 Then
zeros += 1
Else
' Set bit in BitRank
Dim As Integer blockIdx = i Shr 6
Dim As Integer bitposic = i And 63
wm.B(h).bits(blockIdx) = wm.B(h).bits(blockIdx) Or (1ULL Shl bitposic)
End If
Next
' Build cnt array
Dim As Integer runningCount = 0
For i = 0 To Ubound(wm.B(h).bits)
Dim As Integer bitCount = 0
Dim As Ulong bits = wm.B(h).bits(i)
' Count bits in this block
While bits <> 0
If (bits And 1) <> 0 Then bitCount += 1
bits Shr= 1
Wend
wm.B(h).count(i) = runningCount
runningCount += bitCount
Next
' Stable partition - separate 0's and 1's
Dim As Integer zeroValues(zeros - 1)
Dim As Integer oneValues(n - zeros - 1)
Dim As Integer zeroIdx = 0, oneIdx = 0
For i = 0 To n - 1
If GetBitAt(tmpVec(i), wm.height - h - 1) = 0 Then
zeroValues(zeroIdx) = tmpVec(i)
zeroIdx += 1
Else
oneValues(oneIdx) = tmpVec(i)
oneIdx += 1
End If
Next
' Combine arrays back
For i = 0 To zeros - 1
tmpVec(i) = zeroValues(i)
Next
For i = 0 To n - zeros - 1
tmpVec(zeros + i) = oneValues(i)
Next
wm.posic(h) = zeros
Next
End Sub
' kth smallest element in [l, r)
Function Quantile(Byref wm As WaveletMatrix, k As Integer, l As Integer, r As Integer) As Integer
Dim As Integer i, zeros, res = 0
For i = 0 To wm.height - 1
zeros = Rank0(wm.B(i), r) - Rank0(wm.B(i), l)
If zeros > k Then
' Go to left (0) child
l = Rank0(wm.B(i), l)
r = Rank0(wm.B(i), r)
Else
' Go to right (1) child
l = wm.posic(i) + Rank1(wm.B(i), l)
r = wm.posic(i) + Rank1(wm.B(i), r)
k -= zeros
res Or= (1 Shl (wm.height - i - 1))
End If
Next
Return res
End Function
#include "wavelet_matrix.bas"
' Function to find the maximum value in an array
Function findMax(arr() As Integer) As Integer
Dim As Integer max = arr(0)
For i As Integer = 1 To Ubound(arr)
If arr(i) > max Then max = arr(i)
Next
Return max
End Function
' Quicksort implementation for sorting arrays
Sub quickSort(arr() As Integer, low As Integer, high As Integer)
If low >= high Then Exit Sub
Dim As Integer i = low, j = high
Dim As Integer pivot = arr((low + high) \ 2)
Do
While arr(i) < pivot: i += 1: Wend
While arr(j) > pivot: j -= 1: Wend
If i <= j Then
Swap arr(i), arr(j)
i += 1
j -= 1
End If
Loop Until i > j
If low < j Then quickSort(arr(), low, j)
If i < high Then quickSort(arr(), i, high)
End Sub
' Find index of an element in the sorted array using binary search
Function findIdx(sortedArr() As Integer, x As Integer, size As Integer) As Integer
Dim As Integer izda = 0, medio, dcha = size - 1
While izda <= dcha
medio = (izda + dcha) \ 2
If sortedArr(medio) = x Then Return medio
If sortedArr(medio) < x Then
izda = medio + 1
Else
dcha = medio - 1
End If
Wend
Return izda
End Function
Sub main()
' Define the array and queries
Dim As Integer a(5) = {3374, 956, 2114, 3415, 3437}
Dim As Integer i, n = Ubound(a)
' Create a sorted copy with unique values
Dim As Integer sortedA(n-1)
For i = 0 To n - 1
sortedA(i) = a(i)
Next
quickSort(sortedA(), 0, n - 1)
' Count unique values
Dim As Integer uniqueCount = 1
For i = 1 To n - 1
If sortedA(i) <> sortedA(i-1) Then uniqueCount += 1
Next
' Create array with unique values
Dim As Integer uniqueA(uniqueCount-1)
uniqueA(0) = sortedA(0)
Dim As Integer idx = 1
For i = 1 To n - 1
If sortedA(i) <> sortedA(i-1) Then
uniqueA(idx) = sortedA(i)
idx += 1
End If
Next
' Map original values to indices in unique array
Dim As Integer mapped(n-1)
For i = 0 To n - 1
mapped(i) = findIdx(uniqueA(), a(i), uniqueCount)
Next
' Create wavelet matrix
Dim As WaveletMatrix wm
CreateWaveletMatrix(wm, mapped(), n, uniqueCount)
' Process queries
Dim queries(4, 2) As Integer = { _
{2, 2, 1}, {3, 4, 1}, {4, 5, 1}, {1, 2, 2}, {4, 4, 1} }
For i = Lbound(queries) To Ubound(queries)
Dim l As Integer = queries(i, 0) - 1 ' Convert to 0-indexed
Dim r As Integer = queries(i, 1)
Dim k As Integer = queries(i, 2)
Print uniqueA(Quantile(wm, k - 1, l, r))
Next
End Sub
main()
Sleep
- Output:
956 2114 3415 3374 3415
package main
import (
"fmt"
"math/bits"
"sort"
)
// BitRank is a rank data structure for bit vectors
type BitRank struct {
block []uint64
count []int
}
// Resize resizes the bit vector to the given length
func (br *BitRank) Resize(num int) {
br.block = make([]uint64, (num+1)>>6+1)
br.count = make([]int, len(br.block))
}
// Set sets bit at position i
func (br *BitRank) Set(i int, val int) {
if val == 1 {
br.block[i>>6] |= (1 << uint(i&63))
}
}
// Build builds the rank structure
func (br *BitRank) Build() {
for i := 1; i < len(br.block); i++ {
br.count[i] = br.count[i-1] + br.popcountll(br.block[i-1])
}
}
// popcountll counts number of 1's in a 64-bit integer
func (br *BitRank) popcountll(n uint64) int {
return bits.OnesCount64(n)
}
// Rank1 counts number of 1's in [0, i)
func (br *BitRank) Rank1(i int) int {
return br.count[i>>6] +
br.popcountll(br.block[i>>6]&((1<<uint(i&63))-1))
}
// Rank1FromTo counts number of 1's in [i, j)
func (br *BitRank) Rank1FromTo(i, j int) int {
return br.Rank1(j) - br.Rank1(i)
}
// Rank0 counts number of 0's in [0, i)
func (br *BitRank) Rank0(i int) int {
return i - br.Rank1(i)
}
// Rank0FromTo counts number of 0's in [i, j)
func (br *BitRank) Rank0FromTo(i, j int) int {
return br.Rank0(j) - br.Rank0(i)
}
// WaveletMatrix is a wavelet matrix data structure
type WaveletMatrix struct {
height int
B []*BitRank
pos []int
}
// NewWaveletMatrix creates a new wavelet matrix
func NewWaveletMatrix(vec []int, sigma ...int) *WaveletMatrix {
wm := &WaveletMatrix{}
s := 0
if len(sigma) > 0 {
s = sigma[0]
} else {
// Find the maximum element and use that as sigma
for _, v := range vec {
if v > s {
s = v
}
}
s++
}
wm.init(vec, s)
return wm
}
func (wm *WaveletMatrix) init(vec []int, sigma int) {
// Calculate height based on sigma value
if sigma == 1 {
wm.height = 1
} else {
wm.height = 64 - bits.LeadingZeros(uint(sigma-1))
}
wm.B = make([]*BitRank, wm.height)
wm.pos = make([]int, wm.height)
for i := 0; i < wm.height; i++ {
wm.B[i] = &BitRank{}
wm.B[i].Resize(len(vec))
for j := 0; j < len(vec); j++ {
wm.B[i].Set(j, wm.get(vec[j], wm.height-i-1))
}
wm.B[i].Build()
// Stable partition - separate 0's and 1's while preserving order
wm.pos[i] = wm.stablePartition(vec, func(c int) bool {
return wm.get(c, wm.height-i-1) == 0
})
}
}
// stablePartition is equivalent to C++ stable_partition
func (wm *WaveletMatrix) stablePartition(arr []int, predicate func(int) bool) int {
result := make([]int, 0, len(arr))
falseValues := make([]int, 0, len(arr))
for _, item := range arr {
if predicate(item) {
result = append(result, item)
} else {
falseValues = append(falseValues, item)
}
}
partitionPoint := len(result)
result = append(result, falseValues...)
// Update the original array
copy(arr, result)
return partitionPoint
}
// get returns bit at position i from val
func (wm *WaveletMatrix) get(val, i int) int {
return (val >> i) & 1
}
// Rank counts occurrences of val in range [l, r)
func (wm *WaveletMatrix) Rank(val, l, r int) int {
return wm.RankSingle(val, r) - wm.RankSingle(val, l)
}
// RankSingle counts occurrences of val in range [0, i)
func (wm *WaveletMatrix) RankSingle(val, i int) int {
p := 0
for j := 0; j < wm.height; j++ {
if wm.get(val, wm.height-j-1) == 1 {
p = wm.pos[j] + wm.B[j].Rank1(p)
i = wm.pos[j] + wm.B[j].Rank1(i)
} else {
p = wm.B[j].Rank0(p)
i = wm.B[j].Rank0(i)
}
}
return i - p
}
// Quantile returns kth smallest element in [l, r)
func (wm *WaveletMatrix) Quantile(k, l, r int) int {
res := 0
for i := 0; i < wm.height; i++ {
j := wm.B[i].Rank0FromTo(l, r)
if j > k {
l = wm.B[i].Rank0(l)
r = wm.B[i].Rank0(r)
} else {
l = wm.pos[i] + wm.B[i].Rank1(l)
r = wm.pos[i] + wm.B[i].Rank1(r)
k -= j
res |= (1 << (wm.height - i - 1))
}
}
return res
}
// RangeFreq counts elements in [l, r) that are in value range [a, b)
func (wm *WaveletMatrix) RangeFreq(l, r, a, b int) int {
return wm.rangeFreqRecursive(l, r, a, b, 0, 1<<wm.height, 0)
}
func (wm *WaveletMatrix) rangeFreqRecursive(i, j, a, b, l, r, x int) int {
if i == j || r <= a || b <= l {
return 0
}
mid := (l + r) >> 1
if a <= l && r <= b {
return j - i
} else {
left := wm.rangeFreqRecursive(
wm.B[x].Rank0(i),
wm.B[x].Rank0(j),
a, b, l, mid, x+1,
)
right := wm.rangeFreqRecursive(
wm.pos[x]+wm.B[x].Rank1(i),
wm.pos[x]+wm.B[x].Rank1(j),
a, b, mid, r, x+1,
)
return left + right
}
}
// RangeMin finds minimum value in [l, r) within value range [a, b), -1 if not found
func (wm *WaveletMatrix) RangeMin(l, r, a, b int) int {
return wm.rangeMinRecursive(l, r, a, b, 0, 1<<wm.height, 0, 0)
}
func (wm *WaveletMatrix) rangeMinRecursive(i, j, a, b, l, r, x, val int) int {
if i == j || r <= a || b <= l {
return -1
}
if r-l == 1 {
return val
}
mid := (l + r) >> 1
res := wm.rangeMinRecursive(
wm.B[x].Rank0(i),
wm.B[x].Rank0(j),
a, b, l, mid, x+1, val,
)
if res < 0 {
return wm.rangeMinRecursive(
wm.pos[x]+wm.B[x].Rank1(i),
wm.pos[x]+wm.B[x].Rank1(j),
a, b, mid, r, x+1,
val+(1<<(wm.height-x-1)),
)
} else {
return res
}
}
// binary search to find index in sorted array
func find(arr []int, x int) int {
left := 0
right := len(arr)
for left < right {
mid := (left + right) / 2
if arr[mid] < x {
left = mid + 1
} else {
right = mid
}
}
return left
}
func main() {
n := 5
a := []int{3374, 956, 2114, 3415, 3437}
input := make([]int, n)
copy(input, a)
backup := make([]int, n)
copy(backup, a)
// Sort and deduplicate the array
sortedA := make([]int, n)
copy(sortedA, a)
sort.Ints(sortedA)
// Deduplicate
uniqueA := []int{}
for i, val := range sortedA {
if i == 0 || val != sortedA[i-1] {
uniqueA = append(uniqueA, val)
}
}
// Map original values to their indices in the unique array
for i := 0; i < n; i++ {
input[i] = find(uniqueA, backup[i])
}
lrkVector := [][]int{
{2, 2, 1},
{3, 4, 1},
{4, 5, 1},
{1, 2, 2},
{4, 4, 1},
}
wm := NewWaveletMatrix(input)
for _, lrk := range lrkVector {
l, r, k := lrk[0], lrk[1], lrk[2]
l-- // Convert to 0-indexed
fmt.Println(uniqueA[wm.Quantile(k-1, l, r)])
}
}
- Output:
956 2114 3415 3374 3415
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.Predicate;
public class WaveletMatrixDemo {
// BitRank is a rank data structure for bit vectors
static class BitRank {
private long[] block;
private int[] count;
// Resize resizes the bit vector to the given length
public void resize(int num) {
block = new long[((num + 1) >> 6) + 1];
count = new int[block.length];
}
// Set sets bit at position i
public void set(int i, int val) {
if (val == 1) {
block[i >> 6] |= (1L << (i & 63));
}
}
// Build builds the rank structure
public void build() {
for (int i = 1; i < block.length; i++) {
count[i] = count[i - 1] + popcountll(block[i - 1]);
}
}
// popcountll counts number of 1's in a 64-bit integer
private int popcountll(long n) {
return Long.bitCount(n);
}
// Rank1 counts number of 1's in [0, i)
public int rank1(int i) {
return count[i >> 6] + popcountll(block[i >> 6] & ((1L << (i & 63)) - 1));
}
// Rank1FromTo counts number of 1's in [i, j)
public int rank1FromTo(int i, int j) {
return rank1(j) - rank1(i);
}
// Rank0 counts number of 0's in [0, i)
public int rank0(int i) {
return i - rank1(i);
}
// Rank0FromTo counts number of 0's in [i, j)
public int rank0FromTo(int i, int j) {
return rank0(j) - rank0(i);
}
}
// WaveletMatrix is a wavelet matrix data structure
static class WaveletMatrix {
private int height;
private BitRank[] B;
private int[] pos;
// Constructor creates a new wavelet matrix
public WaveletMatrix(int[] vec, int... sigma) {
int s = 0;
if (sigma.length > 0) {
s = sigma[0];
} else {
// Find the maximum element and use that as sigma
for (int v : vec) {
if (v > s) {
s = v;
}
}
s++;
}
init(vec, s);
}
private void init(int[] vec, int sigma) {
// Calculate height based on sigma value
if (sigma == 1) {
height = 1;
} else {
height = 64 - Long.numberOfLeadingZeros(sigma - 1);
}
B = new BitRank[height];
pos = new int[height];
for (int i = 0; i < height; i++) {
B[i] = new BitRank();
B[i].resize(vec.length);
for (int j = 0; j < vec.length; j++) {
B[i].set(j, get(vec[j], height - i - 1));
}
B[i].build();
// Use a final variable to capture the current i value
final int currentLevel = i;
pos[i] = stablePartition(vec, c -> get(c, height - currentLevel - 1) == 0);
}
}
// stablePartition is equivalent to C++ stable_partition
private int stablePartition(int[] arr, Predicate<Integer> predicate) {
List<Integer> result = new ArrayList<>(arr.length);
List<Integer> falseValues = new ArrayList<>(arr.length);
for (int item : arr) {
if (predicate.test(item)) {
result.add(item);
} else {
falseValues.add(item);
}
}
int partitionPoint = result.size();
result.addAll(falseValues);
// Update the original array
for (int i = 0; i < result.size(); i++) {
arr[i] = result.get(i);
}
return partitionPoint;
}
// get returns bit at position i from val
private int get(int val, int i) {
return (val >> i) & 1;
}
// Rank counts occurrences of val in range [l, r)
public int rank(int val, int l, int r) {
return rankSingle(val, r) - rankSingle(val, l);
}
// RankSingle counts occurrences of val in range [0, i)
public int rankSingle(int val, int i) {
int p = 0;
for (int j = 0; j < height; j++) {
if (get(val, height - j - 1) == 1) {
p = pos[j] + B[j].rank1(p);
i = pos[j] + B[j].rank1(i);
} else {
p = B[j].rank0(p);
i = B[j].rank0(i);
}
}
return i - p;
}
// Quantile returns kth smallest element in [l, r)
public int quantile(int k, int l, int r) {
int res = 0;
for (int i = 0; i < height; i++) {
int j = B[i].rank0FromTo(l, r);
if (j > k) {
l = B[i].rank0(l);
r = B[i].rank0(r);
} else {
l = pos[i] + B[i].rank1(l);
r = pos[i] + B[i].rank1(r);
k -= j;
res |= (1 << (height - i - 1));
}
}
return res;
}
// RangeFreq counts elements in [l, r) that are in value range [a, b)
public int rangeFreq(int l, int r, int a, int b) {
return rangeFreqRecursive(l, r, a, b, 0, 1 << height, 0);
}
private int rangeFreqRecursive(int i, int j, int a, int b, int l, int r, int x) {
if (i == j || r <= a || b <= l) {
return 0;
}
int mid = (l + r) >> 1;
if (a <= l && r <= b) {
return j - i;
} else {
int left = rangeFreqRecursive(
B[x].rank0(i),
B[x].rank0(j),
a, b, l, mid, x + 1
);
int right = rangeFreqRecursive(
pos[x] + B[x].rank1(i),
pos[x] + B[x].rank1(j),
a, b, mid, r, x + 1
);
return left + right;
}
}
// RangeMin finds minimum value in [l, r) within value range [a, b), -1 if not found
public int rangeMin(int l, int r, int a, int b) {
return rangeMinRecursive(l, r, a, b, 0, 1 << height, 0, 0);
}
private int rangeMinRecursive(int i, int j, int a, int b, int l, int r, int x, int val) {
if (i == j || r <= a || b <= l) {
return -1;
}
if (r - l == 1) {
return val;
}
int mid = (l + r) >> 1;
int res = rangeMinRecursive(
B[x].rank0(i),
B[x].rank0(j),
a, b, l, mid, x + 1, val
);
if (res < 0) {
return rangeMinRecursive(
pos[x] + B[x].rank1(i),
pos[x] + B[x].rank1(j),
a, b, mid, r, x + 1,
val + (1 << (height - x - 1))
);
} else {
return res;
}
}
}
// binary search to find index in sorted array
private static int find(int[] arr, int x) {
int left = 0;
int right = arr.length;
while (left < right) {
int mid = (left + right) / 2;
if (arr[mid] < x) {
left = mid + 1;
} else {
right = mid;
}
}
return left;
}
public static void main(String[] args) {
int n = 5;
int[] a = {3374, 956, 2114, 3415, 3437};
int[] input = Arrays.copyOf(a, n);
int[] backup = Arrays.copyOf(a, n);
// Sort and deduplicate the array
int[] sortedA = Arrays.copyOf(a, n);
Arrays.sort(sortedA);
// Deduplicate
List<Integer> uniqueAList = new ArrayList<>();
for (int i = 0; i < sortedA.length; i++) {
if (i == 0 || sortedA[i] != sortedA[i - 1]) {
uniqueAList.add(sortedA[i]);
}
}
// Convert List to array
int[] uniqueA = new int[uniqueAList.size()];
for (int i = 0; i < uniqueAList.size(); i++) {
uniqueA[i] = uniqueAList.get(i);
}
// Map original values to their indices in the unique array
for (int i = 0; i < n; i++) {
input[i] = find(uniqueA, backup[i]);
}
int[][] lrkVector = {
{2, 2, 1},
{3, 4, 1},
{4, 5, 1},
{1, 2, 2},
{4, 4, 1}
};
WaveletMatrix wm = new WaveletMatrix(input);
for (int[] lrk : lrkVector) {
int l = lrk[0];
int r = lrk[1];
int k = lrk[2];
l--; // Convert to 0-indexed
System.out.println(uniqueA[wm.quantile(k - 1, l, r)]);
}
}
}
- Output:
956 2114 3415 3374 3415
class BitRank {
constructor() {
this.block = [];
this.count = [];
}
// Resize the bit vector to the given length
resize(num) {
this.block = new Array(Math.floor((num + 1) >> 6) + 1).fill(0n);
this.count = new Array(this.block.length).fill(0);
}
// Set bit at position i
set(i, val) {
this.block[i >> 6] |= (BigInt(val) << BigInt(i & 63));
}
// Build the rank structure
build() {
for (let i = 1; i < this.block.length; i++) {
this.count[i] = this.count[i - 1] + this.popcountll(this.block[i - 1]);
}
}
// Count number of 1's in a 64-bit integer (popcount equivalent)
popcountll(n) {
let count = 0;
while (n) {
count += Number(n & 1n);
n >>= 1n;
}
return count;
}
// Count number of 1's in [0, i)
rank1(i) {
return this.count[i >> 6] +
this.popcountll(this.block[i >> 6] & ((1n << BigInt(i & 63)) - 1n));
}
// Count number of 1's in [i, j)
rank1FromTo(i, j) {
return this.rank1(j) - this.rank1(i);
}
// Count number of 0's in [0, i)
rank0(i) {
return i - this.rank1(i);
}
// Count number of 0's in [i, j)
rank0FromTo(i, j) {
return this.rank0(j) - this.rank0(i);
}
}
class WaveletMatrix {
constructor(vec, sigma) {
if (sigma === undefined) {
// Find the maximum element and use that as sigma
sigma = Math.max(...vec) + 1;
}
this.init(vec, sigma);
}
init(vec, sigma) {
// Calculate height based on sigma value
this.height = (sigma === 1) ? 1 : (64 - Math.clz32(sigma - 1));
this.B = new Array(this.height);
this.pos = new Array(this.height);
for (let i = 0; i < this.height; i++) {
this.B[i] = new BitRank();
this.B[i].resize(vec.length);
for (let j = 0; j < vec.length; j++) {
this.B[i].set(j, this.get(vec[j], this.height - i - 1));
}
this.B[i].build();
// Stable partition - separate 0's and 1's while preserving order
const partition = this.stablePartition(vec, c => !this.get(c, this.height - i - 1));
this.pos[i] = partition;
}
}
// Stable partition implementation (equivalent to C++ stable_partition)
stablePartition(arr, predicate) {
const result = [];
const falseValues = [];
for (const item of arr) {
if (predicate(item)) {
result.push(item);
} else {
falseValues.push(item);
}
}
const partitionPoint = result.length;
result.push(...falseValues);
// Update the original array
for (let i = 0; i < arr.length; i++) {
arr[i] = result[i];
}
return partitionPoint;
}
// Get bit at position i from val
get(val, i) {
return (val >> i) & 1;
}
// Count occurrences of val in range [l, r)
rank(val, l, r) {
if (r === undefined) {
// Single parameter version: [0, l)
return this.rankSingle(val, l);
}
return this.rankSingle(val, r) - this.rankSingle(val, l);
}
// Count occurrences of val in range [0, i)
rankSingle(val, i) {
let p = 0;
for (let j = 0; j < this.height; j++) {
if (this.get(val, this.height - j - 1)) {
p = this.pos[j] + this.B[j].rank1(p);
i = this.pos[j] + this.B[j].rank1(i);
} else {
p = this.B[j].rank0(p);
i = this.B[j].rank0(i);
}
}
return i - p;
}
// kth smallest element in [l, r)
quantile(k, l, r) {
let res = 0;
for (let i = 0; i < this.height; i++) {
const j = this.B[i].rank0FromTo(l, r);
if (j > k) {
l = this.B[i].rank0(l);
r = this.B[i].rank0(r);
} else {
l = this.pos[i] + this.B[i].rank1(l);
r = this.pos[i] + this.B[i].rank1(r);
k -= j;
res |= (1 << (this.height - i - 1));
}
}
return res;
}
// Count elements in [l, r) that are in value range [a, b)
rangefreq(l, r, a, b) {
return this.rangefreqRecursive(l, r, a, b, 0, 1 << this.height, 0);
}
rangefreqRecursive(i, j, a, b, l, r, x) {
if (i === j || r <= a || b <= l) return 0;
const mid = (l + r) >> 1;
if (a <= l && r <= b) {
return j - i;
} else {
const left = this.rangefreqRecursive(
this.B[x].rank0(i),
this.B[x].rank0(j),
a, b, l, mid, x + 1
);
const right = this.rangefreqRecursive(
this.pos[x] + this.B[x].rank1(i),
this.pos[x] + this.B[x].rank1(j),
a, b, mid, r, x + 1
);
return left + right;
}
}
// Find minimum value in [l, r) within value range [a, b), -1 if not found
rangemin(l, r, a, b) {
return this.rangeminRecursive(l, r, a, b, 0, 1 << this.height, 0, 0);
}
rangeminRecursive(i, j, a, b, l, r, x, val) {
if (i === j || r <= a || b <= l) return -1;
if (r - l === 1) return val;
const mid = (l + r) >> 1;
const res = this.rangeminRecursive(
this.B[x].rank0(i),
this.B[x].rank0(j),
a, b, l, mid, x + 1, val
);
if (res < 0) {
return this.rangeminRecursive(
this.pos[x] + this.B[x].rank1(i),
this.pos[x] + this.B[x].rank1(j),
a, b, mid, r, x + 1,
val + (1 << (this.height - x - 1))
);
} else {
return res;
}
}
}
// Main function
function main() {
const n = 5;
const q = 5;
const a = [ 3374, 956, 2114, 3415, 3437 ]
const input = [...a];
const backup = [...a];
// Sort and deduplicate the array
const sortedA = [...a].sort((a, b) => a - b);
const uniqueA = [...new Set(sortedA)];
// Function to find index of an element in the unique array
const find = (x) => {
let left = 0;
let right = uniqueA.length;
while (left < right) {
const mid = Math.floor((left + right) / 2);
if (uniqueA[mid] < x) {
left = mid + 1;
} else {
right = mid;
}
}
return left;
};
// Map original values to their indices in the unique array
for (let i = 0; i < n; i++) {
input[i] = find(backup[i]);
}
const lrk_vector = [[2, 2, 1], [3, 4, 1], [4, 5, 1], [1, 2, 2], [4, 4, 1]];
const wm = new WaveletMatrix(input);
for (const lrk of lrk_vector) {
let [l, r, k] = lrk;
l--; // Convert to 0-indexed
console.log(uniqueA[wm.quantile(k - 1, l, r)]);
}
}
// Execute the main function
main();
- Output:
956 2114 3415 3374 3415
""" rosettacode.org/wiki/Wavelet_Matrix """
""" A block of bits, within an Int vector, and their ranks """
mutable struct BitRank
block::Vector{UInt64}
count::Vector{UInt64}
BitRank(num) = (n = (num + 1) >> 6 + 1; new(zeros(UInt64, n), zeros(UInt64, n)))
end
""" Set bit at position i """
function set!(br::BitRank, i, val)
br.block[i>>6+1] |= (val << (i & 63))
end
""" Build the rank structure """
function build!(br::BitRank)
for i in 2:length(br.block)
br.count[i] = br.count[i-1] + popcountll(br.block[i-1])
end
end
""" Count number of 1's in an integer (popcount equivalent) """
function popcountll(n)
count = 0
while n != 0
count += n & 1
n >>= 1
end
return count
end
""" Count number of 1's in [0, i) """
function rank1(br::BitRank, i)
block_idx = i >> 6 + 1
br.count[block_idx] + popcountll(br.block[block_idx] & ((1 << (i & 63)) - 1))
end
""" Count number of 1's in [i, j) """
rank1FromTo(br::BitRank, i, j) = rank1(br, j) - rank1(br, i)
""" Count number of 0's in [0, i) """
rank0(br::BitRank, i) = i - rank1(br, i)
""" Count number of 0's in [i, j) """
rank0FromTo(br::BitRank, i, j) = rank0(br, j) - rank0(br, i)
""" Wavelet object with bit rankings """
mutable struct WaveletMatrix
height::Int
B::Vector{BitRank}
pos::Vector{Int}
end
function WaveletMatrix(vec::Vector{Int}, sigma = nothing)
isnothing(sigma) && (sigma = maximum(vec) + 1)
len = length(vec)
# Calculate height based on sigma value
height = sigma == 1 ? 1 : 65 - ndigits(sigma - 1, base = 2)
wm = WaveletMatrix(height, BitRank[], Int[])
for i in 1:wm.height
push!(wm.B, BitRank(len))
for j in 1:len
set!(wm.B[i], j - 1, getbit(vec[j], wm.height - i))
end
build!(wm.B[i])
# Stable partition - separate 0's and 1's while preserving order
push!(wm.pos, stable_partition!(vec, c -> getbit(c, wm.height - i) == 0))
end
return wm
end
""" Stable partition implementation """
function stable_partition!(arr::Vector{Int}, predicate)
result = Int[]
false_values = Int[]
for item in arr
if predicate(item)
push!(result, item)
else
push!(false_values, item)
end
end
partition_point = length(result)
append!(result, false_values)
# Update the original array
copyto!(arr, result)
return partition_point
end
""" Get bit at position i from val """
getbit(val, i) = (val >> i) & 1
""" Count occurrences of val in range [l, r) """
function rank(wm::WaveletMatrix, val, l, r = nothing)
isnothing(r) && return rank_single(wm, val, l)
return rank_single(wm, val, r) - rank_single(wm, val, l)
end
""" Count occurrences of val in range [0, i) """
function rank_single(wm::WaveletMatrix, val, i)
p = 0
for j in 1:wm.height
if getbit(val, wm.height - j)
p = wm.pos[j] + rank1(wm.B[j], p)
i = wm.pos[j] + rank1(wm.B[j], i)
else
p = rank0(wm.B[j], p)
i = rank0(wm.B[j], i)
end
end
return i - p
end
""" kth smallest element in [l, r) """
function quantile(wm::WaveletMatrix, k, l, r)
res = 0
for i in 1:wm.height
j = rank0FromTo(wm.B[i], l, r)
if j > k
l = rank0(wm.B[i], l)
r = rank0(wm.B[i], r)
else
l = Int(wm.pos[i] + rank1(wm.B[i], l))
r = Int(wm.pos[i] + rank1(wm.B[i], r))
k -= j
res |= (1 << (wm.height - i))
end
end
return res
end
""" Count elements in [l, r) that are in value range [a, b) """
function rangefreq(wm::WaveletMatrix, l, r, a, b)
return rangefreq_recursive(wm, l, r, a, b, 0, 1 << wm.height, 0)
end
function rangefreq_recursive(wm::WaveletMatrix, i, j, a, b, l, r, x)
if i == j || r <= a || b <= l
return 0
end
mid = (l + r) >> 1
if a <= l && r <= b
return j - i
end
left = rangefreq_recursive(wm, rank0(wm.B[x+1], i), rank0(wm.B[x+1], j), a, b, l, mid, x + 1)
right = rangefreq_recursive(wm, wm.pos[x+1] + rank1(wm.B[x+1], i), wm.pos[x+1] + rank1(wm.B[x+1], j), a, b, mid, r, x + 1)
return left + right
end
""" Find minimum value in [l, r) within value range [a, b), -1 if not found """
function rangemin(wm::WaveletMatrix, l, r, a, b)
rangemin_recursive(wm, l, r, a, b, 0, 1 << wm.height, 0, 0)
end
function rangemin_recursive(wm::WaveletMatrix, i, j, a, b, l, r, x, val)
(i == j || r <= a || b <= l) && return -1
r - l == 1 && return val
mid = (l + r) >> 1
res = rangemin_recursive(wm, rank0(wm.B[x+1], i), rank0(wm.B[x+1], j), a, b, l, mid, x + 1, val)
return res < 0 ?
rangemin_recursive(wm, wm.pos[x+1] + rank1(wm.B[x+1], i), wm.pos[x+1] + rank1(wm.B[x+1], j), a, b, mid, r, x + 1, val + (1 << (wm.height - x - 1))) :
res
end
""" test the wavelet matrix """
function main()
a = [3374, 956, 2114, 3415, 3437]
unique_a = unique!(sort(a))
input = [findfirst(==(x), unique_a) - 1 for x in a]
wm = WaveletMatrix(input)
lrk_vector = [[2, 2, 1], [3, 4, 1], [4, 5, 1], [1, 2, 2], [4, 4, 1]]
for (l, r, k) in lrk_vector
println(unique_a[quantile(wm, k - 1, l - 1, r)+1])
end
end
# Execute the main function
main()
- Output:
Same as the JavaScript example output.
import java.util.ArrayList
class MainKt {
// BitRank is a rank data structure for bit vectors
class BitRank {
private lateinit var block: LongArray
private lateinit var count: IntArray
// Resize resizes the bit vector to the given length
fun resize(num: Int) {
block = LongArray(((num + 1) shr 6) + 1)
count = IntArray(block.size)
}
// Set sets bit at position i
fun set(i: Int, value: Int) {
if (value == 1) {
block[i shr 6] = block[i shr 6] or (1L shl (i and 63))
}
}
// Build builds the rank structure
fun build() {
for (i in 1 until block.size) {
count[i] = count[i - 1] + popcountll(block[i - 1])
}
}
// popcountll counts number of 1's in a 64-bit integer
private fun popcountll(n: Long): Int {
return java.lang.Long.bitCount(n)
}
// Rank1 counts number of 1's in [0, i)
fun rank1(i: Int): Int {
return count[i shr 6] + popcountll(block[i shr 6] and ((1L shl (i and 63)) - 1))
}
// Rank1FromTo counts number of 1's in [i, j)
fun rank1FromTo(i: Int, j: Int): Int {
return rank1(j) - rank1(i)
}
// Rank0 counts number of 0's in [0, i)
fun rank0(i: Int): Int {
return i - rank1(i)
}
// Rank0FromTo counts number of 0's in [i, j)
fun rank0FromTo(i: Int, j: Int): Int {
return rank0(j) - rank0(i)
}
}
// WaveletMatrix is a wavelet matrix data structure
class WaveletMatrix {
private var height: Int = 0
private lateinit var B: Array<BitRank>
private lateinit var pos: IntArray
// Constructor creates a new wavelet matrix
constructor(vec: IntArray, vararg sigma: Int) {
var s = 0
if (sigma.isNotEmpty()) {
s = sigma[0]
} else {
// Find the maximum element and use that as sigma
for (v in vec) {
if (v > s) {
s = v
}
}
s++
}
init(vec, s)
}
private fun init(vec: IntArray, sigma: Int) {
// Calculate height based on sigma value
height = if (sigma == 1) {
1
} else {
64 - java.lang.Long.numberOfLeadingZeros((sigma - 1).toLong())
}
B = Array(height) { BitRank() }
pos = IntArray(height)
for (i in 0 until height) {
B[i].resize(vec.size)
for (j in vec.indices) {
B[i].set(j, get(vec[j], height - i - 1))
}
B[i].build()
pos[i] = stablePartition(vec) { get(it, height - i - 1) == 0 }
}
}
// stablePartition is equivalent to C++ stable_partition
private fun stablePartition(arr: IntArray, predicate: (Int) -> Boolean): Int {
val result = ArrayList<Int>(arr.size)
val falseValues = ArrayList<Int>(arr.size)
for (item in arr) {
if (predicate(item)) {
result.add(item)
} else {
falseValues.add(item)
}
}
val partitionPoint = result.size
result.addAll(falseValues)
// Update the original array
for (i in result.indices) {
arr[i] = result[i]
}
return partitionPoint
}
// get returns bit at position i from val
private fun get(value: Int, i: Int): Int {
return (value shr i) and 1
}
// Rank counts occurrences of val in range [l, r)
fun rank(value: Int, l: Int, r: Int): Int {
return rankSingle(value, r) - rankSingle(value, l)
}
// RankSingle counts occurrences of val in range [0, i)
fun rankSingle(value: Int, i: Int): Int {
var p = 0
var idx = i
for (j in 0 until height) {
if (get(value, height - j - 1) == 1) {
p = pos[j] + B[j].rank1(p)
idx = pos[j] + B[j].rank1(idx)
} else {
p = B[j].rank0(p)
idx = B[j].rank0(idx)
}
}
return idx - p
}
// Quantile returns kth smallest element in [l, r)
fun quantile(k: Int, l: Int, r: Int): Int {
var res = 0
var left = l
var right = r
var kVal = k
for (i in 0 until height) {
val j = B[i].rank0FromTo(left, right)
if (j > kVal) {
left = B[i].rank0(left)
right = B[i].rank0(right)
} else {
left = pos[i] + B[i].rank1(left)
right = pos[i] + B[i].rank1(right)
kVal -= j
res = res or (1 shl (height - i - 1))
}
}
return res
}
// RangeFreq counts elements in [l, r) that are in value range [a, b)
fun rangeFreq(l: Int, r: Int, a: Int, b: Int): Int {
return rangeFreqRecursive(l, r, a, b, 0, 1 shl height, 0)
}
private fun rangeFreqRecursive(i: Int, j: Int, a: Int, b: Int, l: Int, r: Int, x: Int): Int {
if (i == j || r <= a || b <= l) {
return 0
}
val mid = (l + r) shr 1
return if (a <= l && r <= b) {
j - i
} else {
val left = rangeFreqRecursive(
B[x].rank0(i),
B[x].rank0(j),
a, b, l, mid, x + 1
)
val right = rangeFreqRecursive(
pos[x] + B[x].rank1(i),
pos[x] + B[x].rank1(j),
a, b, mid, r, x + 1
)
left + right
}
}
// RangeMin finds minimum value in [l, r) within value range [a, b), -1 if not found
fun rangeMin(l: Int, r: Int, a: Int, b: Int): Int {
return rangeMinRecursive(l, r, a, b, 0, 1 shl height, 0, 0)
}
private fun rangeMinRecursive(i: Int, j: Int, a: Int, b: Int, l: Int, r: Int, x: Int, value: Int): Int {
if (i == j || r <= a || b <= l) {
return -1
}
if (r - l == 1) {
return value
}
val mid = (l + r) shr 1
val res = rangeMinRecursive(
B[x].rank0(i),
B[x].rank0(j),
a, b, l, mid, x + 1, value
)
return if (res < 0) {
rangeMinRecursive(
pos[x] + B[x].rank1(i),
pos[x] + B[x].rank1(j),
a, b, mid, r, x + 1,
value + (1 shl (height - x - 1))
)
} else {
res
}
}
}
companion object {
// binary search to find index in sorted array
private fun find(arr: IntArray, x: Int): Int {
var left = 0
var right = arr.size
while (left < right) {
val mid = (left + right) / 2
if (arr[mid] < x) {
left = mid + 1
} else {
right = mid
}
}
return left
}
@JvmStatic
fun main(args: Array<String>) {
val n = 5
val a = intArrayOf(3374, 956, 2114, 3415, 3437)
val input = a.copyOf(n)
val backup = a.copyOf(n)
// Sort and deduplicate the array
val sortedA = a.copyOf(n).also { it.sort() }
// Deduplicate
val uniqueAList = ArrayList<Int>()
for (i in sortedA.indices) {
if (i == 0 || sortedA[i] != sortedA[i - 1]) {
uniqueAList.add(sortedA[i])
}
}
// Convert List to array
val uniqueA = IntArray(uniqueAList.size)
for (i in uniqueAList.indices) {
uniqueA[i] = uniqueAList[i]
}
// Map original values to their indices in the unique array
for (i in 0 until n) {
input[i] = find(uniqueA, backup[i])
}
val lrkVector = arrayOf(
intArrayOf(2, 2, 1),
intArrayOf(3, 4, 1),
intArrayOf(4, 5, 1),
intArrayOf(1, 2, 2),
intArrayOf(4, 4, 1)
)
val wm = WaveletMatrix(input)
for (lrk in lrkVector) {
val l = lrk[0] - 1 // Convert to 0-indexed
val r = lrk[1]
val k = lrk[2]
println(uniqueA[wm.quantile(k - 1, l, r)])
}
}
}
}
- Output:
956 2114 3415 3374 3415
(* BitRank is a rank data structure for bit vectors *)
module BitRank = struct
type t = {
mutable block: int64 array;
mutable count: int array;
}
let create () = { block = [||]; count = [||] }
(* Resize resizes the bit vector to the given length *)
let resize br num =
let block_size = ((num + 1) lsr 6) + 1 in
br.block <- Array.make block_size 0L;
br.count <- Array.make block_size 0
(* Set sets bit at position i *)
let set br i val_ =
if val_ = 1 then
let idx = i lsr 6 in
let mask = Int64.shift_left 1L (i land 63) in
br.block.(idx) <- Int64.logor br.block.(idx) mask
(* popcountll counts number of 1's in a 64-bit integer *)
let popcountll n =
let rec count_bits n acc =
if n = 0L then acc
else count_bits (Int64.logand n (Int64.sub n 1L)) (acc + 1)
in
count_bits n 0
(* Build builds the rank structure *)
let build br =
for i = 1 to Array.length br.block - 1 do
br.count.(i) <- br.count.(i - 1) + popcountll br.block.(i - 1)
done
(* Rank1 counts number of 1's in [0, i) *)
let rank1 br i =
let idx = i lsr 6 in
let mask = Int64.pred (Int64.shift_left 1L (i land 63)) in
br.count.(idx) + popcountll (Int64.logand br.block.(idx) mask)
(* Rank1FromTo counts number of 1's in [i, j) *)
let rank1_from_to br i j = rank1 br j - rank1 br i
(* Rank0 counts number of 0's in [0, i) *)
let rank0 br i = i - rank1 br i
(* Rank0FromTo counts number of 0's in [i, j) *)
let rank0_from_to br i j = rank0 br j - rank0 br i
end
(* WaveletMatrix is a wavelet matrix data structure *)
module WaveletMatrix = struct
type t = {
mutable height: int;
mutable b: BitRank.t array;
mutable pos: int array;
}
(* get returns bit at position i from val *)
let get val_ i = (val_ lsr i) land 1
(* stablePartition is equivalent to C++ stable_partition *)
let stable_partition arr predicate =
let result = ref [] in
let false_values = ref [] in
Array.iter (fun item ->
if predicate item then
result := item :: !result
else
false_values := item :: !false_values
) arr;
let result_list = List.rev !result in
let partition_point = List.length result_list in
let combined = result_list @ (List.rev !false_values) in
(* Update the original array *)
List.iteri (fun i item -> arr.(i) <- item) combined;
partition_point
(* Helper function to count leading zeros in an int *)
let leading_zeros n =
if n = 0 then 64 else
let rec count_zeros n bit acc =
if bit < 0 then acc
else if n land (1 lsl bit) <> 0 then acc
else count_zeros n (bit - 1) (acc + 1)
in
count_zeros n 63 0
(* Initialize wavelet matrix *)
let init vec sigma =
let wm = { height = 0; b = [||]; pos = [||] } in
(* Calculate height based on sigma value *)
wm.height <-
if sigma = 1 then 1
else 64 - leading_zeros (sigma - 1);
wm.b <- Array.init wm.height (fun _ -> BitRank.create ());
wm.pos <- Array.make wm.height 0;
for i = 0 to wm.height - 1 do
BitRank.resize wm.b.(i) (Array.length vec);
for j = 0 to Array.length vec - 1 do
BitRank.set wm.b.(i) j (get vec.(j) (wm.height - i - 1))
done;
BitRank.build wm.b.(i);
(* Use a closure to capture the current i value *)
let current_level = i in
wm.pos.(i) <- stable_partition vec (fun c -> get c (wm.height - current_level - 1) = 0)
done;
wm
(* Create a new wavelet matrix *)
let create vec sigma_opt =
let sigma = match sigma_opt with
| Some s -> s
| None ->
(* Find the maximum element and use that as sigma *)
let max_val = Array.fold_left max 0 vec in
max_val + 1
in
init vec sigma
(* RankSingle counts occurrences of val in range [0, i) *)
let rank_single wm val_ i =
let p = ref 0 in
let i_ref = ref i in
for j = 0 to wm.height - 1 do
if get val_ (wm.height - j - 1) = 1 then begin
p := wm.pos.(j) + BitRank.rank1 wm.b.(j) !p;
i_ref := wm.pos.(j) + BitRank.rank1 wm.b.(j) !i_ref
end else begin
p := BitRank.rank0 wm.b.(j) !p;
i_ref := BitRank.rank0 wm.b.(j) !i_ref
end
done;
!i_ref - !p
(* Rank counts occurrences of val in range [l, r) *)
let rank wm val_ l r = rank_single wm val_ r - rank_single wm val_ l
(* Quantile returns kth smallest element in [l, r) *)
let quantile wm k l r =
let res = ref 0 in
let l_ref = ref l in
let r_ref = ref r in
let k_ref = ref k in
for i = 0 to wm.height - 1 do
let j = BitRank.rank0_from_to wm.b.(i) !l_ref !r_ref in
if j > !k_ref then begin
l_ref := BitRank.rank0 wm.b.(i) !l_ref;
r_ref := BitRank.rank0 wm.b.(i) !r_ref
end else begin
l_ref := wm.pos.(i) + BitRank.rank1 wm.b.(i) !l_ref;
r_ref := wm.pos.(i) + BitRank.rank1 wm.b.(i) !r_ref;
k_ref := !k_ref - j;
res := !res lor (1 lsl (wm.height - i - 1))
end
done;
!res
(* RangeFreq counts elements in [l, r) that are in value range [a, b) *)
let rec range_freq_recursive wm i j a b l r x =
if i = j || r <= a || b <= l then
0
else if a <= l && r <= b then
j - i
else begin
let mid = (l + r) lsr 1 in
let left = range_freq_recursive wm
(BitRank.rank0 wm.b.(x) i)
(BitRank.rank0 wm.b.(x) j)
a b l mid (x + 1)
in
let right = range_freq_recursive wm
(wm.pos.(x) + BitRank.rank1 wm.b.(x) i)
(wm.pos.(x) + BitRank.rank1 wm.b.(x) j)
a b mid r (x + 1)
in
left + right
end
let range_freq wm l r a b =
range_freq_recursive wm l r a b 0 (1 lsl wm.height) 0
(* RangeMin finds minimum value in [l, r) within value range [a, b), -1 if not found *)
let rec range_min_recursive wm i j a b l r x val_ =
if i = j || r <= a || b <= l then
-1
else if r - l = 1 then
val_
else begin
let mid = (l + r) lsr 1 in
let res = range_min_recursive wm
(BitRank.rank0 wm.b.(x) i)
(BitRank.rank0 wm.b.(x) j)
a b l mid (x + 1) val_
in
if res < 0 then
range_min_recursive wm
(wm.pos.(x) + BitRank.rank1 wm.b.(x) i)
(wm.pos.(x) + BitRank.rank1 wm.b.(x) j)
a b mid r (x + 1)
(val_ + (1 lsl (wm.height - x - 1)))
else
res
end
let range_min wm l r a b =
range_min_recursive wm l r a b 0 (1 lsl wm.height) 0 0
end
(* binary search to find index in sorted array *)
let find arr x =
let rec binary_search left right =
if left >= right then left
else
let mid = (left + right) / 2 in
if arr.(mid) < x then
binary_search (mid + 1) right
else
binary_search left mid
in
binary_search 0 (Array.length arr)
let () =
let n = 5 in
let a = [|3374; 956; 2114; 3415; 3437|] in
let input = Array.copy a in
let backup = Array.copy a in
(* Sort and deduplicate the array *)
let sorted_a = Array.copy a in
Array.sort compare sorted_a;
(* Deduplicate *)
let unique_a_list = ref [] in
for i = 0 to Array.length sorted_a - 1 do
if i = 0 || sorted_a.(i) <> sorted_a.(i - 1) then
unique_a_list := sorted_a.(i) :: !unique_a_list
done;
(* Convert List to array *)
let unique_a = Array.of_list (List.rev !unique_a_list) in
(* Map original values to their indices in the unique array *)
for i = 0 to n - 1 do
input.(i) <- find unique_a backup.(i)
done;
let lrk_vector = [|
[|2; 2; 1|];
[|3; 4; 1|];
[|4; 5; 1|];
[|1; 2; 2|];
[|4; 4; 1|]
|] in
let wm = WaveletMatrix.create input None in
Array.iter (fun lrk ->
let l = lrk.(0) in
let r = lrk.(1) in
let k = lrk.(2) in
let l = l - 1 in (* Convert to 0-indexed *)
Printf.printf "%d\n" unique_a.(WaveletMatrix.quantile wm (k - 1) l r)
) lrk_vector
- Output:
956 2114 3415 3374 3415
// BitRank "class" for rank operations
enum BITS, COUNT
type BitRank(object br)
return sequence(br)
and sequence(br[BITS])
and sequence(br[COUNT])
end type
// WaveletMatrix "class"
enum B, // Array of BitRanks
POS, // position array
HEIGHT // Height of the matrix
type WaveletMatrix(object wm)
return sequence(wm)
and length(wm)=3
and sequence(wm[B])
and sequence(wm[POS])
and integer(wm[HEIGHT])
end type
// Get i'th bit of v
function GetBitAt(integer v, i)
return and_bits(floor(v/power(2,i)),1)
end function
// Count number of 1's in [0, i)
function Rank1(BitRank br, integer i)
integer blockIndex = floor(i/64),
bitposic = and_bits(i,63),
cnt = 0
// Add cnt from previous blocks
if blockIndex > 0 then cnt = br[COUNT][blockIndex] end if
// Count bits in current block
for j=1 to bitposic do
if and_bits(br[BITS][blockIndex+1],power(2,j-1)) then cnt += 1 end if
end for
return cnt
end function
// Count number of 0's in [0, i)
function Rank0(BitRank br, integer i)
return i - Rank1(br, i)
end function
// Initialize a BitRank structure
function InitBitRank(integer size)
integer numBlocks = floor((size+63)/64)
BitRank br = {repeat(0,numBlocks), -- [BITS]
repeat(0,numBlocks)} -- [COUNT]
return br
end function
// Initialize the WaveletMatrix with a vector
function CreateWaveletMatrix(sequence vec, integer maxVal=-1)
if maxVal=-1 then maxVal = max(vec)+1 end if
integer height = iff(maxVal=1?1:64-(32-length(int_to_bits(maxVal)))),
n = length(vec)
WaveletMatrix wm = {repeat(0,height), -- B
repeat(0,height), -- POS
height} -- HEIGHT
// Create a copy of the input vector for manipulation
sequence tmpVec = deep_copy(vec)
for h=1 to height do
wm[B][h] = InitBitRank(n)
// Count zeros
integer zeros = 0,
bit = height-h
for i=1 to n do
integer vi = tmpVec[i]-1
if GetBitAt(vi, bit) = 0 then
zeros += 1
else
// Set bit in BitRank
integer blockIdx = floor((i-1)/64),
bitposic = and_bits((i-1),63)
wm[B][h][BITS][blockIdx+1] ||= power(2,bitposic)
end if
end for
// Build cnt array
integer runningCount = 0
for i,bits in wm[B][h][BITS] do
wm[B][h][COUNT][i] = runningCount
runningCount += count_bits(bits)
end for
// Stable partition - separate 0's and 1's
sequence zeroValues = repeat(0,zeros),
oneValues = repeat(0,n-zeros)
integer zeroIdx = 1, oneIdx = 1
for i,ti in tmpVec do
if GetBitAt(ti-1, bit) = 0 then
zeroValues[zeroIdx] = ti
zeroIdx += 1
else
oneValues[oneIdx] = ti
oneIdx += 1
end if
end for
// Combine arrays back
zeroValues &= oneValues
tmpVec = zeroValues
wm[POS][h] = zeros
end for
return wm
end function
// kth smallest element in [l, r)
function Quantile(WaveletMatrix wm, integer k, l, r)
integer res = 0, h = wm[HEIGHT]
for i=1 to h do
BitRank b = wm[B][i]
integer nl = Rank0(b, l),
nr = Rank0(b, r),
zeros = nr - nl
if zeros > k then
// Go to left (0) child
l = nl
r = nr
else
// Go to right (1) child
integer p = wm[POS][i]
l = p + l-nl
r = p + r-nr
k -= zeros
res ||= power(2,h-i)
end if
end for
return res
end function
procedure main()
sequence a = {3374, 956, 2114, 3415, 3437},
uniqa = unique(a),
mapped = apply(true,binary_search,{a,{uniqa}})
WaveletMatrix wm = CreateWaveletMatrix(mapped, length(uniqa))
for query in {{2,2,1},{3,4,1},{4,5,1},{1,2,2},{4,4,1}} do
integer {l,r,k} = query
?uniqa[Quantile(wm, k-1, l-1, r)+1]
end for
end procedure
main()
- Output:
956 2114 3415 3374 3415
import math
class BitRank:
def __init__(self):
self.block = []
self.count = []
def resize(self, num):
self.block = [0] * (math.floor((num + 1) >> 6) + 1)
self.count = [0] * len(self.block)
def set(self, i, val):
self.block[i >> 6] |= (val << (i & 63))
def build(self):
for i in range(1, len(self.block)):
self.count[i] = self.count[i - 1] + \
self.popcountll(self.block[i - 1])
def popcountll(self, n):
count = 0
while n:
count += n & 1
n >>= 1
return count
def rank1(self, i):
return self.count[i >> 6] + self.popcountll(self.block[i >> 6] & ((1 << (i & 63)) - 1))
def rank1FromTo(self, i, j):
return self.rank1(j) - self.rank1(i)
def rank0(self, i):
return i - self.rank1(i)
def rank0FromTo(self, i, j):
return self.rank0(j) - self.rank0(i)
class WaveletMatrix:
def __init__(self, vec, sigma=None):
if sigma is None:
sigma = max(vec) + 1
self.init(vec, sigma)
def init(self, vec, sigma):
# Calculate height based on sigma value
self.height = 1 if sigma == 1 else (64 - (sigma - 1).bit_length() + 1)
self.B = [None] * self.height
self.pos = [0] * self.height
for i in range(self.height):
self.B[i] = BitRank()
self.B[i].resize(len(vec))
for j in range(len(vec)):
self.B[i].set(j, self.get(vec[j], self.height - i - 1))
self.B[i].build()
# Stable partition
partition = self.stable_partition(
vec, lambda c: not self.get(c, self.height - i - 1))
self.pos[i] = partition
def stable_partition(self, arr, predicate):
result = []
false_values = []
for item in arr:
if predicate(item):
result.append(item)
else:
false_values.append(item)
partition_point = len(result)
result.extend(false_values)
# Update the original array
for i in range(len(arr)):
arr[i] = result[i]
return partition_point
def get(self, val, i):
return (val >> i) & 1
def rank(self, val, l, r=None):
if r is None:
return self.rank_single(val, l)
return self.rank_single(val, r) - self.rank_single(val, l)
def rank_single(self, val, i):
p = 0
for j in range(self.height):
if self.get(val, self.height - j - 1):
p = self.pos[j] + self.B[j].rank1(p)
i = self.pos[j] + self.B[j].rank1(i)
else:
p = self.B[j].rank0(p)
i = self.B[j].rank0(i)
return i - p
def quantile(self, k, l, r):
res = 0
for i in range(self.height):
j = self.B[i].rank0FromTo(l, r)
if j > k:
l = self.B[i].rank0(l)
r = self.B[i].rank0(r)
else:
l = self.pos[i] + self.B[i].rank1(l)
r = self.pos[i] + self.B[i].rank1(r)
k -= j
res |= (1 << (self.height - i - 1))
return res
def rangefreq(self, l, r, a, b):
return self.rangefreq_recursive(l, r, a, b, 0, 1 << self.height, 0)
def rangefreq_recursive(self, i, j, a, b, l, r, x):
if i == j or r <= a or b <= l:
return 0
mid = (l + r) >> 1
if a <= l and r <= b:
return j - i
left = self.rangefreq_recursive(
self.B[x].rank0(i),
self.B[x].rank0(j),
a, b, l, mid, x + 1
)
right = self.rangefreq_recursive(
self.pos[x] + self.B[x].rank1(i),
self.pos[x] + self.B[x].rank1(j),
a, b, mid, r, x + 1
)
return left + right
def rangemin(self, l, r, a, b):
return self.rangemin_recursive(l, r, a, b, 0, 1 << self.height, 0, 0)
def rangemin_recursive(self, i, j, a, b, l, r, x, val):
if i == j or r <= a or b <= l:
return -1
if r - l == 1:
return val
mid = (l + r) >> 1
res = self.rangemin_recursive(
self.B[x].rank0(i),
self.B[x].rank0(j),
a, b, l, mid, x + 1, val
)
if res < 0:
return self.rangemin_recursive(
self.pos[x] + self.B[x].rank1(i),
self.pos[x] + self.B[x].rank1(j),
a, b, mid, r, x + 1,
val + (1 << (self.height - x - 1))
)
return res
def main():
n = 5
a = [3374, 956, 2114, 3415, 3437]
input_arr = a.copy()
backup = a.copy()
# Sort and deduplicate the array
sorted_a = sorted(a)
unique_a = list(dict.fromkeys(sorted_a))
def find(x):
left, right = 0, len(unique_a)
while left < right:
mid = (left + right) // 2
if unique_a[mid] < x:
left = mid + 1
else:
right = mid
return left
# Map original values to their indices in the unique array
for i in range(n):
input_arr[i] = find(backup[i])
lrk_vector = [[2, 2, 1], [3, 4, 1], [4, 5, 1], [1, 2, 2], [4, 4, 1]]
wm = WaveletMatrix(input_arr)
for l, r, k in lrk_vector:
l -= 1 # Convert to 0-indexed
print(unique_a[wm.quantile(k - 1, l, r)])
if __name__ == "__main__":
main()
- Output:
Same as JavaScript entry.
EnableExplicit
Structure BitRank
Bits.q[64] ; Array to store bits
Count.i[64] ; Array to store counts
NumBlocks.i ; Number of blocks
EndStructure
Structure WaveletMatrix
B.BitRank[32] ; Array of BitRanks
Posic.i[32] ; Position array
Height.i ; Height of the matrix
EndStructure
; Get bit at position i from value
Procedure.i GetBitAt(value.i, bitpos.i)
ProcedureReturn (value >> bitpos) & 1
EndProcedure
; Count number of 1's in [0, i)
Procedure.i Rank1(*br.BitRank, i.i)
If i <= 0 : ProcedureReturn 0 : EndIf
Protected blockIndex.i = i >> 6
Protected bitPos.i = i & 63
Protected count.i = 0, j.i
; Add count from previous blocks
If blockIndex > 0 : count = *br\Count[blockIndex - 1] : EndIf
; Count bits in current block
For j.i = 0 To bitPos - 1
If (*br\Bits[blockIndex] & (1 << j)) <> 0: count + 1 : EndIf
Next
ProcedureReturn count
EndProcedure
; Count number of 0's in [0, i)
Procedure.i Rank0(*br.BitRank, i.i)
ProcedureReturn i - Rank1(*br, i)
EndProcedure
Procedure InitBitRank(*br.BitRank, size.i)
Protected i.i
*br\NumBlocks = (size + 63) >> 6
; Initialize arrays
For i = 0 To 63
If i < *br\NumBlocks
*br\Bits[i] = 0
*br\Count[i] = 0
EndIf
Next
EndProcedure
; Initialize the WaveletMatrix with a vector
Procedure CreateWaveletMatrix(*wm.WaveletMatrix, Array vec(1), n.i, maxVal.i)
Protected.i i, level, zeros, bitpos, blockIdx, runningCount, bitCount, bits
Protected.i zeroIdx, oneIdx, value
; Calculate height based on maximum value
*wm\Height = 0
If maxVal > 0 : *wm\Height = Len(Bin(maxVal)) : EndIf
If *wm\Height < 1 : *wm\Height = 1 : EndIf
; Create a copy of the input vector for manipulation
Dim tmpVec(n-1)
For i = 0 To n - 1
tmpVec(i) = vec(i)
Next
For level = 0 To *wm\Height - 1
If level < 32 ; Initialize BitRank for this level
InitBitRank(@*wm\B[level], n)
; Count zeros
zeros = 0
For i = 0 To n - 1
If GetBitAt(tmpVec(i), *wm\Height - level - 1) = 0
zeros + 1
Else
; Set bit in BitRank
blockIdx = i >> 6
bitPos = i & 63
If blockIdx < 64 : *wm\B[level]\Bits[blockIdx] | (1 << bitPos) : EndIf
EndIf
Next
; Build count array
runningCount = 0
For i = 0 To *wm\B[level]\NumBlocks
If i < 64
bitCount = 0
bits = *wm\B[level]\Bits[i]
; Count bits in this block
While bits <> 0
If (bits & 1) <> 0 : bitCount + 1 :EndIf
bits >> 1
Wend
*wm\B[level]\Count[i] = runningCount
runningCount + bitCount
EndIf
Next
; Stable partition - separate 0's and 1's
Dim zeroValues(zeros-1)
Dim oneValues(n-zeros-1)
zeroIdx = 0
oneIdx = 0
For i = 0 To n - 1
value = tmpVec(i)
If GetBitAt(value, *wm\Height - level - 1) = 0
zeroValues(zeroIdx) = value
zeroIdx + 1
Else
oneValues(oneIdx) = value
oneIdx + 1
EndIf
Next
; Combine arrays back
For i = 0 To zeros - 1
tmpVec(i) = zeroValues(i)
Next
For i = 0 To n - zeros - 1
tmpVec(zeros + i) = oneValues(i)
Next
; Store position
*wm\Posic[level] = zeros
EndIf
Next
EndProcedure
; kth smallest element in [l, r)
Procedure.i Quantile(*wm.WaveletMatrix, k.i, l.i, r.i)
Protected level.i, zeros.i
Protected res.i = 0
For level = 0 To *wm\Height - 1
If level < 32
zeros = Rank0(@*wm\B[level], r) - Rank0(@*wm\B[level], l)
If zeros > k
; Go to left (0) child
l = Rank0(@*wm\B[level], l)
r = Rank0(@*wm\B[level], r)
Else
; Go to right (1) child
l = *wm\Posic[level] + Rank1(@*wm\B[level], l)
r = *wm\Posic[level] + Rank1(@*wm\B[level], r)
k - zeros
res | (1 << (*wm\Height - level - 1))
EndIf
EndIf
Next
ProcedureReturn res
EndProcedure
; Main program
If OpenConsole()
; Define the array
Dim a.i(4)
a(0) = 3374
a(1) = 956
a(2) = 2114
a(3) = 3415
a(4) = 3437
Define.i i, n, l, r, k
Define.i uniqueCount, idx, izda, dcha, medio
n = 5
; Create a sorted copy with unique values
Dim sortedA.i(n-1)
For i = 0 To n - 1
sortedA(i) = a(i)
Next
; Sort the array
SortArray(sortedA(), #PB_Sort_Ascending)
; Count unique values
uniqueCount = 1
For i.i = 1 To n - 1
If sortedA(i) <> sortedA(i-1) : uniqueCount + 1 : EndIf
Next
; Create array with unique values
Dim uniqueA(uniqueCount-1)
uniqueA(0) = sortedA(0)
idx = 1
For i = 1 To n - 1
If sortedA(i) <> sortedA(i-1)
uniqueA(idx) = sortedA(i)
idx + 1
EndIf
Next
; Map original values to indices in unique array
Dim mapped(n-1)
For i = 0 To n - 1
; Find index of a(i) in uniqueA using binary search
izda = 0
dcha = uniqueCount - 1
While izda <= dcha
medio = (izda + dcha) / 2
If uniqueA(medio) = a(i)
mapped(i) = medio
Break
ElseIf uniqueA(medio) < a(i)
izda = medio + 1
Else
dcha = medio - 1
EndIf
Wend
Next
; Create wavelet matrix
Define wm.WaveletMatrix
CreateWaveletMatrix(@wm, mapped(), n, uniqueCount)
; Process queries
Dim queries.i(4, 2)
queries(0, 0) = 2 : queries(0, 1) = 2 : queries(0, 2) = 1
queries(1, 0) = 3 : queries(1, 1) = 4 : queries(1, 2) = 1
queries(2, 0) = 4 : queries(2, 1) = 5 : queries(2, 2) = 1
queries(3, 0) = 1 : queries(3, 1) = 2 : queries(3, 2) = 2
queries(4, 0) = 4 : queries(4, 1) = 4 : queries(4, 2) = 1
For i = 0 To 4
l = queries(i, 0) - 1 ; Convert to 0-indexed
r = queries(i, 1)
k = queries(i, 2)
PrintN(Str(uniqueA(Quantile(@wm, k - 1, l, r))))
Next
PrintN(#CRLF$ + "Press ENTER to exit"): Input()
CloseConsole()
EndIf
- Output:
Siame as FreeBASIC entry.
# 20250625 Raku programming solution
class BitRank {
has array[uint64] $.block is rw;
has array[uint] $.count is rw;
method Resize(Int $num) {
my $size = ($num + 1) +> 6 + 1;
self.block = array[uint64].new: 0 xx $size;
self.count = array[uint].new: 0 xx $size;
}
method SetOne(Int $i, Int $val) {
if $val == 1 { $.block[$i +> 6] +|= (1 +< ($i +& 63)) }
}
method popcountll(uint64 $n is copy) returns Int {
loop ( my $count = 0 ; $n > 0 ; $count++ ) { $n +&= $n - 1 }
return $count
}
method Rank1(Int $i) returns Int {
$.count[$i +> 6] + self.popcountll($.block[$i +> 6] +& ((1 +< ($i +& 63)) - 1));
}
method Rank1FromTo(Int $i, Int $j) returns Int {
return self.Rank1($j) - self.Rank1($i);
}
method Rank0(Int $i) returns Int { return $i - self.Rank1($i) }
method Rank0FromTo(Int $i, Int $j) returns Int {
return self.Rank0($j) - self.Rank0($i)
}
}
class WaveletMatrix {
has Int $.height is rw;
has BitRank @.B is rw;
has Int @.pos is rw;
method new(@vec, Int :$sigma?) { self.bless(:@vec, :$sigma) }
submethod BUILD(:@vec, Int :$sigma?) {
my $s = $sigma // (@vec.max + 1);
self.height = $s <= 1 ?? 1 !! ($s - 1).base(2).chars;
self.B = BitRank.new xx self.height;
self.pos = 0 xx self.height;
for ^self.height -> $i {
my $br = BitRank.new;
$br.Resize(@vec.elems);
for ^@vec.elems -> $j {
$br.SetOne($j, self.get(@vec[$j], self.height - $i - 1));
}
self.B[$i] = $br;
self.pos[$i] = self.stablePartition(@vec, -> $c { self.get($c, self.height - $i - 1) == 0 });
}
}
method stablePartition(@arr, &predicate) returns Int {
my $partition = @arr.classify: { &predicate($_) ?? 'true' !! 'false' };
@arr = do given $partition { |.<true>, |.<false> }
return $partition.<true>.elems
}
method get(Int $val, Int $i) returns Int { return ($val +> $i) +& 1 }
method Rank(Int $val, Int $l, Int $r) returns Int {
return self.RankSingle($val, $r) - self.RankSingle($val, $l);
}
method RankSingle(Int $val, Int $i) returns Int {
my ($p, $idx) = 0, $i;
for 0 ..^ $.height -> $j {
if self.get($val, $.height - $j - 1) == 1 {
$p = @.pos[$j] + @.B[$j].Rank1($p);
$idx = @.pos[$j] + @.B[$j].Rank1($idx);
} else {
$p = @.B[$j].Rank0($p);
$idx = @.B[$j].Rank0($idx);
}
}
return $idx - $p;
}
method Quantile(Int $k, Int $l, Int $r) returns Int {
my ($res, $left, $right, $kth) = 0, $l, $r, $k;
for 0 ..^ $.height -> $i {
my $j = @.B[$i].Rank0FromTo($left, $right);
if $j > $kth {
$left = @.B[$i].Rank0($left);
$right = @.B[$i].Rank0($right);
} else {
$left = @.pos[$i] + @.B[$i].Rank1($left);
$right = @.pos[$i] + @.B[$i].Rank1($right);
$kth -= $j;
$res +|= (1 +< ($.height - $i - 1));
}
}
return $res
}
method RangeFreq(Int $l, Int $r, Int $a, Int $b) returns Int {
return self.rangeFreqRecursive($l, $r, $a, $b, 0, 1 +< $.height, 0)
}
method rangeFreqRecursive(Int $i, Int $j, Int $a, Int $b, Int $l, Int $r, Int $x) returns Int {
if $i == $j || $r <= $a || $b <= $l { return 0 }
my $mid = ($l + $r) +> 1;
if $a <= $l && $r <= $b {
return $j - $i;
} else {
my $left = self.rangeFreqRecursive( @.B[$x].Rank0($i),
@.B[$x].Rank0($j),
$a, $b, $l, $mid, $x + 1);
my $right = self.rangeFreqRecursive( @.pos[$x] + @.B[$x].Rank1($i),
@.pos[$x] + @.B[$x].Rank1($j),
$a, $b, $mid, $r, $x + 1);
return $left + $right
}
}
method rangeMinRecursive(Int $i, Int $j, Int $a, Int $b, Int $l, Int $r, Int $x, Int $val) returns Int {
return -1 if $i == $j || $r <= $a || $b <= $l;
return $val if $r - $l == 1;
my $mid = ($l + $r) +> 1;
my $res = self.rangeMinRecursive( @.B[$x].Rank0($i),
@.B[$x].Rank0($j),
$a, $b, $l, $mid, $x + 1, $val);
if $res < 0 {
return self.rangeMinRecursive( @.pos[$x] + @.B[$x].Rank1($i),
@.pos[$x] + @.B[$x].Rank1($j),
$a, $b, $mid, $r, $x + 1,
$val + (1 +< ($.height - $x - 1)))
} else {
return $res
}
}
}
sub find(@arr, Int $x) returns Int {
my ($left, $right) = 0, @arr.elems;
while $left < $right {
my $mid = ($left + $right) div 2;
@arr[$mid] < $x ?? ( $left = $mid + 1 ) !! ( $right = $mid );
}
return $left
}
sub MAIN() {
my $n = 5;
my @a = 3374, 956, 2114, 3415, 3437;
my @input = my @backup = @a;
my @uniqueA = @a.sort.unique;
for ^$n -> $i { @input[$i] = find(@uniqueA, @backup[$i]) }
my @lrkVector = [ [2, 2, 1], [3, 4, 1], [4, 5, 1], [1, 2, 2], [4, 4, 1], ];
my $wm = WaveletMatrix.new(@input);
for @lrkVector -> [$l, $r, $k] {
my $l_idx = $l - 1; # Convert to 0-indexed
say @uniqueA[$wm.Quantile($k - 1, $l_idx, $r)];
}
}
You may Attempt This Online!
#[derive(Clone)]
struct BitRank {
block: Vec<u64>,
count: Vec<i32>,
}
impl BitRank {
// Resize resizes the bit vector to the given length
fn resize(&mut self, num: usize) {
// Fix the potential overflow in the calculation
self.block = vec![0; ((num + 63) / 64) + 1]; // Safer way to calculate number of u64s needed
self.count = vec![0; self.block.len()];
}
// Set sets bit at position i
fn set(&mut self, i: usize, val: i32) {
if val == 1 {
self.block[i >> 6] |= 1u64 << (i & 63);
}
}
// Build builds the rank structure
fn build(&mut self) {
for i in 1..self.block.len() {
self.count[i] = self.count[i - 1] + self.popcountll(self.block[i - 1]);
}
}
// popcountll counts number of 1's in a 64-bit integer
fn popcountll(&self, n: u64) -> i32 {
n.count_ones() as i32
}
// Rank1 counts number of 1's in [0, i)
fn rank1(&self, i: usize) -> i32 {
if i == 0 {
return 0;
}
let block_idx = i >> 6;
let bit_pos = i & 63;
if block_idx >= self.block.len() {
return self.count[self.count.len() - 1];
}
let mask = if bit_pos == 0 { 0 } else { (1u64 << bit_pos) - 1 };
self.count[block_idx] + self.popcountll(self.block[block_idx] & mask)
}
// Rank1FromTo counts number of 1's in [i, j)
fn rank1_from_to(&self, i: usize, j: usize) -> i32 {
self.rank1(j) - self.rank1(i)
}
// Rank0 counts number of 0's in [0, i)
fn rank0(&self, i: usize) -> i32 {
i as i32 - self.rank1(i)
}
// Rank0FromTo counts number of 0's in [i, j)
fn rank0_from_to(&self, i: usize, j: usize) -> i32 {
self.rank0(j) - self.rank0(i)
}
}
// Helper function to get bit at position i from val
fn get_bit(val: i32, i: i32) -> i32 {
if i < 0 || i >= 32 {
return 0; // Safe default for out of range bit positions
}
(val >> i) & 1
}
// WaveletMatrix is a wavelet matrix data structure
struct WaveletMatrix {
height: i32,
b: Vec<BitRank>,
pos: Vec<i32>,
}
impl WaveletMatrix {
// Create a new wavelet matrix
fn new(vec: &[i32], sigma: Option<i32>) -> Self {
let mut wm = WaveletMatrix {
height: 0,
b: Vec::new(),
pos: Vec::new(),
};
let s = match sigma {
Some(s) => s,
None => {
// Find the maximum element and use that as sigma
let mut max_val = 0;
for &v in vec {
if v > max_val {
max_val = v;
}
}
max_val + 1
}
};
wm.init(vec, s);
wm
}
fn init(&mut self, vec: &[i32], sigma: i32) {
// Calculate height based on sigma value
if sigma <= 1 {
self.height = 1;
} else {
// Safely calculate the height to avoid overflow
self.height = 32 - (sigma - 1).leading_zeros() as i32;
}
self.b = Vec::with_capacity(self.height as usize);
for _ in 0..self.height {
self.b.push(BitRank { block: Vec::new(), count: Vec::new() });
}
self.pos = vec![0; self.height as usize];
let mut temp_vec = vec.to_vec();
for i in 0..self.height as usize {
self.b[i].resize(vec.len());
for j in 0..vec.len() {
// Using the standalone get_bit function to avoid borrowing self twice
let bit_val = get_bit(temp_vec[j], self.height - i as i32 - 1);
self.b[i].set(j, bit_val);
}
self.b[i].build();
// Stable partition - separate 0's and 1's while preserving order
let height_i = self.height - i as i32 - 1; // Calculate once to avoid borrow issues
self.pos[i] = self.stable_partition(&mut temp_vec, move |c| {
get_bit(c, height_i) == 0
}) as i32;
}
}
// stable_partition is equivalent to C++ stable_partition
fn stable_partition<F>(&self, arr: &mut Vec<i32>, predicate: F) -> usize
where F: Fn(i32) -> bool {
let mut result = Vec::with_capacity(arr.len());
let mut false_values = Vec::with_capacity(arr.len());
for &item in arr.iter() {
if predicate(item) {
result.push(item);
} else {
false_values.push(item);
}
}
let partition_point = result.len();
result.extend(false_values);
// Update the original array
arr.copy_from_slice(&result);
partition_point
}
// Rank counts occurrences of val in range [l, r)
fn rank(&self, val: i32, l: usize, r: usize) -> i32 {
self.rank_single(val, r) - self.rank_single(val, l)
}
// RankSingle counts occurrences of val in range [0, i)
fn rank_single(&self, val: i32, i: usize) -> i32 {
let mut p = 0;
let mut i = i as i32;
for j in 0..self.height as usize {
if get_bit(val, self.height - j as i32 - 1) == 1 {
p = self.pos[j] + self.b[j].rank1(p as usize);
i = self.pos[j] + self.b[j].rank1(i as usize);
} else {
p = self.b[j].rank0(p as usize);
i = self.b[j].rank0(i as usize);
}
}
i - p
}
// Quantile returns kth smallest element in [l, r)
fn quantile(&self, k: i32, l: usize, r: usize) -> i32 {
let mut res = 0;
let mut k = k;
let mut l = l;
let mut r = r;
for i in 0..self.height as usize {
let j = self.b[i].rank0_from_to(l, r);
if j > k {
l = self.b[i].rank0(l) as usize;
r = self.b[i].rank0(r) as usize;
} else {
l = (self.pos[i] + self.b[i].rank1(l)) as usize;
r = (self.pos[i] + self.b[i].rank1(r)) as usize;
k -= j;
res |= 1 << (self.height - i as i32 - 1);
}
}
res
}
// RangeFreq counts elements in [l, r) that are in value range [a, b)
fn range_freq(&self, l: usize, r: usize, a: i32, b: i32) -> i32 {
// Use a safer calculation for the maximum value
let max_val = 1i32.checked_shl(self.height as u32).unwrap_or(0);
self.range_freq_recursive(l, r, a, b, 0, max_val, 0)
}
fn range_freq_recursive(&self, i: usize, j: usize, a: i32, b: i32, l: i32, r: i32, x: usize) -> i32 {
if i == j || r <= a || b <= l {
return 0;
}
let mid = (l + r) >> 1;
if a <= l && r <= b {
return (j - i) as i32;
} else {
let left = self.range_freq_recursive(
self.b[x].rank0(i) as usize,
self.b[x].rank0(j) as usize,
a, b, l, mid, x + 1,
);
let right = self.range_freq_recursive(
(self.pos[x] + self.b[x].rank1(i)) as usize,
(self.pos[x] + self.b[x].rank1(j)) as usize,
a, b, mid, r, x + 1,
);
return left + right;
}
}
// RangeMin finds minimum value in [l, r) within value range [a, b), -1 if not found
fn range_min(&self, l: usize, r: usize, a: i32, b: i32) -> i32 {
// Use a safer calculation for the maximum value
let max_val = 1i32.checked_shl(self.height as u32).unwrap_or(0);
self.range_min_recursive(l, r, a, b, 0, max_val, 0, 0)
}
fn range_min_recursive(&self, i: usize, j: usize, a: i32, b: i32, l: i32, r: i32, x: usize, val: i32) -> i32 {
if i == j || r <= a || b <= l {
return -1;
}
if r - l == 1 {
return val;
}
let mid = (l + r) >> 1;
let res = self.range_min_recursive(
self.b[x].rank0(i) as usize,
self.b[x].rank0(j) as usize,
a, b, l, mid, x + 1, val,
);
if res < 0 {
return self.range_min_recursive(
(self.pos[x] + self.b[x].rank1(i)) as usize,
(self.pos[x] + self.b[x].rank1(j)) as usize,
a, b, mid, r, x + 1,
val + (1 << (self.height - x as i32 - 1)),
);
} else {
return res;
}
}
}
// binary search to find index in sorted array
fn find(arr: &[i32], x: i32) -> usize {
let mut left = 0;
let mut right = arr.len();
while left < right {
let mid = (left + right) / 2;
if arr[mid] < x {
left = mid + 1;
} else {
right = mid;
}
}
left
}
fn main() {
let n = 5;
let a = vec![3374, 956, 2114, 3415, 3437];
let mut input = a.clone();
let backup = a.clone();
// Sort and deduplicate the array
let mut sorted_a = a.clone();
sorted_a.sort();
// Deduplicate
let mut unique_a = Vec::new();
for i in 0..sorted_a.len() {
if i == 0 || sorted_a[i] != sorted_a[i-1] {
unique_a.push(sorted_a[i]);
}
}
// Map original values to their indices in the unique array
for i in 0..n {
input[i] = find(&unique_a, backup[i]) as i32;
}
let lrk_vector = vec![
vec![2, 2, 1],
vec![3, 4, 1],
vec![4, 5, 1],
vec![1, 2, 2],
vec![4, 4, 1],
];
let wm = WaveletMatrix::new(&input, None);
for lrk in lrk_vector {
let l = lrk[0] - 1; // Convert to 0-indexed
let r = lrk[1];
let k = lrk[2];
println!("{}", unique_a[wm.quantile(k - 1, l as usize, r as usize) as usize]);
}
}
- Output:
956 2114 3415 3374 3415
import Foundation
class WaveletMatrixDemo {
// BitRank is a rank data structure for bit vectors
class BitRank {
private var block: [UInt64] = []
private var count: [Int] = []
// Resize resizes the bit vector to the given length
func resize(_ num: Int) {
block = Array(repeating: 0, count: ((num + 1) >> 6) + 1)
count = Array(repeating: 0, count: block.count)
}
// Set sets bit at position i
func set(_ i: Int, _ val: Int) {
if val == 1 {
block[i >> 6] |= (1 << (i & 63))
}
}
// Build builds the rank structure
func build() {
for i in 1..<block.count {
count[i] = count[i - 1] + popcountll(block[i - 1])
}
}
// popcountll counts number of 1's in a 64-bit integer
private func popcountll(_ n: UInt64) -> Int {
return n.nonzeroBitCount
}
// Rank1 counts number of 1's in [0, i)
func rank1(_ i: Int) -> Int {
return count[i >> 6] + popcountll(block[i >> 6] & ((1 << (i & 63)) - 1))
}
// Rank1FromTo counts number of 1's in [i, j)
func rank1FromTo(_ i: Int, _ j: Int) -> Int {
return rank1(j) - rank1(i)
}
// Rank0 counts number of 0's in [0, i)
func rank0(_ i: Int) -> Int {
return i - rank1(i)
}
// Rank0FromTo counts number of 0's in [i, j)
func rank0FromTo(_ i: Int, _ j: Int) -> Int {
return rank0(j) - rank0(i)
}
}
// WaveletMatrix is a wavelet matrix data structure
class WaveletMatrix {
private var height: Int = 0
private var B: [BitRank] = []
private var pos: [Int] = []
// Constructor creates a new wavelet matrix
init(_ vec: [Int], _ sigma: Int...) {
var s = 0
if !sigma.isEmpty {
s = sigma[0]
} else {
// Find the maximum element and use that as sigma
for v in vec {
if v > s {
s = v
}
}
s += 1
}
initialize(vec, s)
}
private func initialize(_ vec: [Int], _ sigma: Int) {
// Calculate height based on sigma value
if sigma == 1 {
height = 1
} else {
height = 64 - sigma.leadingZeroBitCount - 1
}
B = Array(repeating: BitRank(), count: height)
pos = Array(repeating: 0, count: height)
var vecCopy = vec
for i in 0..<height {
B[i] = BitRank()
B[i].resize(vecCopy.count)
for j in 0..<vecCopy.count {
B[i].set(j, get(vecCopy[j], height - i - 1))
}
B[i].build()
// Use stablePartition to sort the array
pos[i] = stablePartition(&vecCopy) { get($0, height - i - 1) == 0 }
}
}
// stablePartition is equivalent to C++ stable_partition
private func stablePartition<T>(_ arr: inout [T], _ predicate: (T) -> Bool) -> Int {
var result: [T] = []
var falseValues: [T] = []
for item in arr {
if predicate(item) {
result.append(item)
} else {
falseValues.append(item)
}
}
let partitionPoint = result.count
result.append(contentsOf: falseValues)
// Update the original array
arr = result
return partitionPoint
}
// get returns bit at position i from val
private func get(_ val: Int, _ i: Int) -> Int {
return (val >> i) & 1
}
// Rank counts occurrences of val in range [l, r)
func rank(_ val: Int, _ l: Int, _ r: Int) -> Int {
return rankSingle(val, r) - rankSingle(val, l)
}
// RankSingle counts occurrences of val in range [0, i)
func rankSingle(_ val: Int, _ i: Int) -> Int {
var p = 0
var idx = i
for j in 0..<height {
if get(val, height - j - 1) == 1 {
p = pos[j] + B[j].rank1(p)
idx = pos[j] + B[j].rank1(idx)
} else {
p = B[j].rank0(p)
idx = B[j].rank0(idx)
}
}
return idx - p
}
// Quantile returns kth smallest element in [l, r)
func quantile(_ k: Int, _ l: Int, _ r: Int) -> Int {
var res = 0
var left = l
var right = r
var kValue = k
for i in 0..<height {
let j = B[i].rank0FromTo(left, right)
if j > kValue {
left = B[i].rank0(left)
right = B[i].rank0(right)
} else {
left = pos[i] + B[i].rank1(left)
right = pos[i] + B[i].rank1(right)
kValue -= j
res |= (1 << (height - i - 1))
}
}
return res
}
// RangeFreq counts elements in [l, r) that are in value range [a, b)
func rangeFreq(_ l: Int, _ r: Int, _ a: Int, _ b: Int) -> Int {
return rangeFreqRecursive(l, r, a, b, 0, 1 << height, 0)
}
private func rangeFreqRecursive(_ i: Int, _ j: Int, _ a: Int, _ b: Int, _ l: Int, _ r: Int, _ x: Int) -> Int {
if i == j || r <= a || b <= l {
return 0
}
let mid = (l + r) >> 1
if a <= l && r <= b {
return j - i
} else {
let left = rangeFreqRecursive(
B[x].rank0(i),
B[x].rank0(j),
a, b, l, mid, x + 1
)
let right = rangeFreqRecursive(
pos[x] + B[x].rank1(i),
pos[x] + B[x].rank1(j),
a, b, mid, r, x + 1
)
return left + right
}
}
// RangeMin finds minimum value in [l, r) within value range [a, b), -1 if not found
func rangeMin(_ l: Int, _ r: Int, _ a: Int, _ b: Int) -> Int {
return rangeMinRecursive(l, r, a, b, 0, 1 << height, 0, 0)
}
private func rangeMinRecursive(_ i: Int, _ j: Int, _ a: Int, _ b: Int, _ l: Int, _ r: Int, _ x: Int, _ val: Int) -> Int {
if i == j || r <= a || b <= l {
return -1
}
if r - l == 1 {
return val
}
let mid = (l + r) >> 1
let res = rangeMinRecursive(
B[x].rank0(i),
B[x].rank0(j),
a, b, l, mid, x + 1, val
)
if res < 0 {
return rangeMinRecursive(
pos[x] + B[x].rank1(i),
pos[x] + B[x].rank1(j),
a, b, mid, r, x + 1,
val + (1 << (height - x - 1))
)
} else {
return res
}
}
}
// binary search to find index in sorted array
static func find(_ arr: [Int], _ x: Int) -> Int {
var left = 0
var right = arr.count
while left < right {
let mid = (left + right) / 2
if arr[mid] < x {
left = mid + 1
} else {
right = mid
}
}
return left
}
static func main() {
let n = 5
let a = [3374, 956, 2114, 3415, 3437]
var input = Array(a[0..<n])
let backup = Array(a[0..<n])
// Sort and deduplicate the array
var sortedA = Array(a[0..<n])
sortedA.sort()
// Deduplicate
var uniqueAList: [Int] = []
for i in 0..<sortedA.count {
if i == 0 || sortedA[i] != sortedA[i - 1] {
uniqueAList.append(sortedA[i])
}
}
let uniqueA = uniqueAList
// Map original values to their indices in the unique array
for i in 0..<n {
input[i] = find(uniqueA, backup[i])
}
let lrkVector = [
[2, 2, 1],
[3, 4, 1],
[4, 5, 1],
[1, 2, 2],
[4, 4, 1]
]
let wm = WaveletMatrix(input)
for lrk in lrkVector {
var l = lrk[0]
let r = lrk[1]
let k = lrk[2]
l -= 1 // Convert to 0-indexed
print(uniqueA[wm.quantile(k - 1, l, r)])
}
}
}
// Run the demo
WaveletMatrixDemo.main()
- Output:
956 2114 956 3374 3415
import "./long" for ULong
import "./math" for Nums
import "./fmt" for Conv
import "./seq" for Lst
// Returns the number of leading zeros in a 32 bit binary number.
var Clz32 = Fn.new { |d|
var b = Conv.itoa(d, 2)
return 32 - b.trimStart("0").count
}
class BitRank {
construct new() {
_block = []
_count = []
}
// Resize the bit vector to the given length.
resize(num) {
var len = ((num + 1) >> 6).floor + 1
_block = List.filled(len, null)
for (i in 0...len) _block[i] = ULong.zero
_count = List.filled(len, 0)
}
// Set bit at position i.
set(i, val) {
_block[i >> 6] = _block[i >> 6] | (ULong.new(val) << ULong.new(i & 63))
}
// Build the rank structure.
build() {
for (i in 1..._block.count) _count[i] = count[i-1] + popcountll(_block[i-1])
}
// Count number of 1's in a 64-bit integer (popcount equivalent).
popcountll(n) {
var count = 0
while (n != ULong.zero) {
count = count + (n & ULong.one).toSmall
n = n >> ULong.one
}
return count
}
// Count number of 1's in [0, i).
rank1(i) {
return _count[i >> 6] + popcountll(_block[i >> 6] & ((ULong.one << ULong.new(i & 63)) - ULong.one))
}
// Count number of 1's in [i, j).
rank1FromTo(i, j) { rank1(j) - rank1(i) }
// Count number of 0's in [0, i).
rank0(i) { i - rank1(i) }
// Count number of 0's in [i, j).
rank0FromTo(i, j) { rank0(j) - rank0(i) }
}
class WaveletMatrix {
construct new(vec, sigma) {
if (!sigma) {
// Find the maximum element and use that as sigma.
sigma = Nums.max(vec) + 1
}
init(vec, sigma)
}
init(vec, sigma) {
// Calculate height based on sigma value.
_height = (sigma == 1) ? 1 : (64 - Clz32.call(sigma - 1))
_B = List.filled(_height, 0)
_pos = List.filled(_height, 0)
for (i in 0..._height) {
_B[i] = BitRank.new()
_B[i].resize(vec.count)
for (j in 0...vec.count) {
_B[i].set(j, get(vec[j], _height - i - 1))
}
_B[i].build()
// Stable partition - separate 0's and 1's while preserving order.
var partition = stablePartition(vec) { |c| get(c, _height - i - 1) == 0 }
_pos[i] = partition
}
}
// Stable partition implementation.
stablePartition(arr, predicate) {
var result = []
var falseValues = []
for (item in arr) {
if (predicate.call(item)) {
result.add(item)
} else {
falseValues.add(item)
}
}
var partitionPoint = result.count
result.addAll(falseValues)
// Update the original array.
for (i in 0...arr.count) arr[i] = result[i]
return partitionPoint
}
// Get bit at position i from val.
get(val, i) {
return (val >> i) & 1
}
// Count occurrences of val in range [l, r).
rank(val, l, r) {
if (!r) {
// Single parameter version: [0, l).
return rankSingle(val, l)
}
return rankSingle(val, r) - rankSingle(val, l)
}
// Count occurrences of val in range [0, i).
rankSingle(val, i) {
var p = 0
for (j in 0..._height) {
if (get(val, _height - j - 1) != 0) {
p = _pos[j] + _B[j].rank1(p)
i = _pos[j] + _B[j].rank1(i)
} else {
p = _B[j].rank0(p)
i = _B[j].rank0(i)
}
}
return i - p
}
// kth smallest element in [l, r)
quantile(k, l, r) {
var res = 0
for (i in 0..._height) {
var j = _B[i].rank0FromTo(l, r)
if (j > k) {
l = _B[i].rank0(l)
r = _B[i].rank0(r)
} else {
l = _pos[i] + _B[i].rank1(l)
r = _pos[i] + _B[i].rank1(r)
k = k - j
res = res | (1 << (_height - i - 1))
}
}
return res
}
// Count elements in [l, r) that are in value range [a, b).
rangefreq(l, r, a, b) {
return rangefreqRecursive(l, r, a, b, 0, 1 << _height, 0)
}
rangefreqRecursive(i, j, a, b, l, r, x) {
if (i == j || r <= a || b <= l) return 0
var mid = (l + r) >> 1
if (a <= l && r <= b) {
return j - i
} else {
var left = rangefreqRecursive(
_B[x].rank0(i),
_B[x].rank0(j),
a, b, l, mid, x + 1
)
var right = rangefreqRecursive(
_pos[x] + _B[x].rank1(i),
_pos[x] + _B[x].rank1(j),
a, b, mid, r, x + 1
)
return left + right
}
}
// Find minimum value in [l, r) within value range [a, b), -1 if not found.
rangemin(l, r, a, b) {
return rangeminRecursive(l, r, a, b, 0, 1 << _height, 0, 0)
}
rangeminRecursive(i, j, a, b, l, r, x, val) {
if (i == j || r <= a || b <= l) return -1
if (r - l == 1) return val
var mid = (l + r) >> 1
var res = rangeminRecursive(
_B[x].rank0(i),
_B[x].rank0(j),
a, b, l, mid, x + 1, val
)
if (res < 0) {
return this.rangeminRecursive(
_pos[x] + _B[x].rank1(i),
_pos[x] + _B[x].rank1(j),
a, b, mid, r, x + 1,
val + (1 << (_height - x - 1))
)
} else {
return res
}
}
}
var n = 5
var q = 5
var a = [3374, 956, 2114, 3415, 3437]
var input = a.toList
var backup = a.toList
// Sort and deduplicate the array.
var sortedA = a.toList.sort()
var uniqueA = Lst.distinct(sortedA)
// Function to find index of an element in the unique array.
var find = Fn.new { |x|
var left = 0
var right = uniqueA.count
while (left < right) {
var mid = ((left + right)/2).floor
if (uniqueA[mid] < x) {
left = mid + 1
} else {
right = mid
}
}
return left
}
// Map original values to their indices in the unique array.
for (i in 0...n) {
input[i] = find.call(backup[i])
}
var lrkVector = [[2, 2, 1], [3, 4, 1], [4, 5, 1], [1, 2, 2], [4, 4, 1]]
var wm = WaveletMatrix.new(input, null)
for (lrk in lrkVector) {
var l = lrk[0]
var r = lrk[1]
var k = lrk[2]
l = l - 1 // convert to 0-indexed
System.print(uniqueA[wm.quantile(k - 1, l, r)])
}
- Output:
956 2114 3415 3374 3415
const std = @import("std");
const BitRank = struct {
block: std.ArrayList(u64),
count: std.ArrayList(i32),
const Self = @This();
fn init(allocator: std.mem.Allocator) Self {
return .{
.block = std.ArrayList(u64).init(allocator),
.count = std.ArrayList(i32).init(allocator),
};
}
fn deinit(self: *Self) void {
self.block.deinit();
self.count.deinit();
}
fn clone(self: Self, allocator: std.mem.Allocator) !Self {
var result = Self.init(allocator);
try result.block.appendSlice(self.block.items);
try result.count.appendSlice(self.count.items);
return result;
}
// Resize resizes the bit vector to the given length
fn resize(self: *Self, num: usize) !void {
try self.block.resize(((num + 63) / 64) + 1);
try self.count.resize(self.block.items.len);
for (self.block.items) |*b| {
b.* = 0;
}
for (self.count.items) |*c| {
c.* = 0;
}
}
// Set sets bit at position i
fn set(self: *Self, i: usize, val: i32) void {
if (val == 1) {
self.block.items[i >> 6] |= @as(u64, 1) << @as(u6, @intCast(i & 63));
}
}
// Build builds the rank structure
fn build(self: *Self) void {
for (1..self.block.items.len) |i| {
self.count.items[i] = self.count.items[i - 1] + self.popcountll(self.block.items[i - 1]);
}
}
// popcountll counts number of 1's in a 64-bit integer
fn popcountll(_: *const Self, n: u64) i32 {
return @as(i32, @intCast(@popCount(n)));
}
// Rank1 counts number of 1's in [0, i)
fn rank1(self: *const Self, i: usize) i32 {
if (i == 0) {
return 0;
}
const block_idx = i >> 6;
const bit_pos = i & 63;
if (block_idx >= self.block.items.len) {
return self.count.items[self.count.items.len - 1];
}
const mask = if (bit_pos == 0) 0 else (@as(u64, 1) << @as(u6, @intCast(bit_pos))) - 1;
return self.count.items[block_idx] + self.popcountll(self.block.items[block_idx] & mask);
}
// Rank1FromTo counts number of 1's in [i, j)
fn rank1FromTo(self: *const Self, i: usize, j: usize) i32 {
return self.rank1(j) - self.rank1(i);
}
// Rank0 counts number of 0's in [0, i)
fn rank0(self: *const Self, i: usize) i32 {
return @as(i32, @intCast(i)) - self.rank1(i);
}
// Rank0FromTo counts number of 0's in [i, j)
fn rank0FromTo(self: *const Self, i: usize, j: usize) i32 {
return self.rank0(j) - self.rank0(i);
}
};
// Helper function to get bit at position i from val
fn getBit(val: i32, i: i32) i32 {
if (i < 0 or i >= 32) {
return 0; // Safe default for out of range bit positions
}
return (val >> @as(u5, @intCast(@as(u32, @intCast(i))))) & 1;
}
// WaveletMatrix is a wavelet matrix data structure
const WaveletMatrix = struct {
height: i32,
b: std.ArrayList(BitRank),
pos: std.ArrayList(i32),
allocator: std.mem.Allocator,
const Self = @This();
fn init(allocator: std.mem.Allocator) Self {
return .{
.height = 0,
.b = std.ArrayList(BitRank).init(allocator),
.pos = std.ArrayList(i32).init(allocator),
.allocator = allocator,
};
}
fn deinit(self: *Self) void {
for (self.b.items) |*bit_rank| {
bit_rank.deinit();
}
self.b.deinit();
self.pos.deinit();
}
// Create a new wavelet matrix
fn new(allocator: std.mem.Allocator, vec: []const i32, sigma: ?i32) !Self {
var wm = Self.init(allocator);
const s = if (sigma) |s| s else blk: {
// Find the maximum element and use that as sigma
var max_val: i32 = 0;
for (vec) |v| {
if (v > max_val) {
max_val = v;
}
}
break :blk max_val + 1;
};
try wm.initWithVector(vec, s);
return wm;
}
fn initWithVector(self: *Self, vec: []const i32, sigma: i32) !void {
// Calculate height based on sigma value
if (sigma <= 1) {
self.height = 1;
} else {
// Safely calculate the height to avoid overflow
self.height = 32 - @as(i32, @intCast(@clz(@as(u32, @intCast(sigma - 1)))));
}
try self.b.resize(@as(usize, @intCast(self.height)));
for (self.b.items) |*bit_rank| {
bit_rank.* = BitRank.init(self.allocator);
}
try self.pos.resize(@as(usize, @intCast(self.height)));
const temp_vec = try self.allocator.alloc(i32, vec.len);
defer self.allocator.free(temp_vec);
@memcpy(temp_vec, vec);
for (0..@as(usize, @intCast(self.height))) |i| {
try self.b.items[i].resize(vec.len);
for (0..vec.len) |j| {
const bit_val = getBit(temp_vec[j], self.height - @as(i32, @intCast(i)) - 1);
self.b.items[i].set(j, bit_val);
}
self.b.items[i].build();
// Stable partition - separate 0's and 1's while preserving order
const height_i = self.height - @as(i32, @intCast(i)) - 1;
self.pos.items[i] = @as(i32, @intCast(self.stablePartition(temp_vec, height_i)));
}
}
// stablePartition is equivalent to C++ stable_partition
fn stablePartition(self: *Self, arr: []i32, height_i: i32) usize {
var result = std.ArrayList(i32).init(self.allocator);
defer result.deinit();
var false_values = std.ArrayList(i32).init(self.allocator);
defer false_values.deinit();
for (arr) |item| {
if (getBit(item, height_i) == 0) {
result.append(item) catch unreachable;
} else {
false_values.append(item) catch unreachable;
}
}
const partition_point = result.items.len;
result.appendSlice(false_values.items) catch unreachable;
// Update the original array
@memcpy(arr, result.items);
return partition_point;
}
// Rank counts occurrences of val in range [l, r)
fn rank(self: *const Self, val: i32, l: usize, r: usize) i32 {
return self.rankSingle(val, r) - self.rankSingle(val, l);
}
// RankSingle counts occurrences of val in range [0, i)
fn rankSingle(self: *const Self, val: i32, i: usize) i32 {
var p: i32 = 0;
var i_val: i32 = @intCast(i);
for (0..@as(usize, @intCast(self.height))) |j| {
if (getBit(val, self.height - @as(i32, @intCast(j)) - 1) == 1) {
p = self.pos.items[j] + self.b.items[j].rank1(@as(usize, @intCast(p)));
i_val = self.pos.items[j] + self.b.items[j].rank1(@as(usize, @intCast(i_val)));
} else {
p = self.b.items[j].rank0(@as(usize, @intCast(p)));
i_val = self.b.items[j].rank0(@as(usize, @intCast(i_val)));
}
}
return i_val - p;
}
// Quantile returns kth smallest element in [l, r)
fn quantile(self: *const Self, k: i32, l: usize, r: usize) i32 {
var res: i32 = 0;
var k_val = k;
var l_val = l;
var r_val = r;
for (0..@as(usize, @intCast(self.height))) |i| {
const j = self.b.items[i].rank0FromTo(l_val, r_val);
if (j > k_val) {
l_val = @as(usize, @intCast(self.b.items[i].rank0(l_val)));
r_val = @as(usize, @intCast(self.b.items[i].rank0(r_val)));
} else {
l_val = @as(usize, @intCast(self.pos.items[i] + self.b.items[i].rank1(l_val)));
r_val = @as(usize, @intCast(self.pos.items[i] + self.b.items[i].rank1(r_val)));
k_val -= j;
res |= @as(i32, 1) << @as(u5, @intCast(self.height - @as(i32, @intCast(i)) - 1));
}
}
return res;
}
// RangeFreq counts elements in [l, r) that are in value range [a, b)
fn rangeFreq(self: *const Self, l: usize, r: usize, a: i32, b: i32) i32 {
// Use a safer calculation for the maximum value
const max_val = if (self.height >= 31) std.math.maxInt(i32) else (@as(i32, 1) << @as(u5, @intCast(self.height)));
return self.rangeFreqRecursive(l, r, a, b, 0, max_val, 0);
}
fn rangeFreqRecursive(self: *const Self, i: usize, j: usize, a: i32, b: i32, l: i32, r: i32, x: usize) i32 {
if (i == j or r <= a or b <= l) {
return 0;
}
const mid = (l + r) >> 1;
if (a <= l and r <= b) {
return @as(i32, @intCast(j - i));
} else {
const left = self.rangeFreqRecursive(
@as(usize, @intCast(self.b.items[x].rank0(i))),
@as(usize, @intCast(self.b.items[x].rank0(j))),
a, b, l, mid, x + 1
);
const right = self.rangeFreqRecursive(
@as(usize, @intCast(self.pos.items[x] + self.b.items[x].rank1(i))),
@as(usize, @intCast(self.pos.items[x] + self.b.items[x].rank1(j))),
a, b, mid, r, x + 1
);
return left + right;
}
}
// RangeMin finds minimum value in [l, r) within value range [a, b), -1 if not found
fn rangeMin(self: *const Self, l: usize, r: usize, a: i32, b: i32) i32 {
// Use a safer calculation for the maximum value
const max_val = if (self.height >= 31) std.math.maxInt(i32) else (@as(i32, 1) << @as(u5, @intCast(self.height)));
return self.rangeMinRecursive(l, r, a, b, 0, max_val, 0, 0);
}
fn rangeMinRecursive(self: *const Self, i: usize, j: usize, a: i32, b: i32, l: i32, r: i32, x: usize, val: i32) i32 {
if (i == j or r <= a or b <= l) {
return -1;
}
if (r - l == 1) {
return val;
}
const mid = (l + r) >> 1;
const res = self.rangeMinRecursive(
@as(usize, @intCast(self.b.items[x].rank0(i))),
@as(usize, @intCast(self.b.items[x].rank0(j))),
a, b, l, mid, x + 1, val
);
if (res < 0) {
return self.rangeMinRecursive(
@as(usize, @intCast(self.pos.items[x] + self.b.items[x].rank1(i))),
@as(usize, @intCast(self.pos.items[x] + self.b.items[x].rank1(j))),
a, b, mid, r, x + 1,
val + (@as(i32, 1) << @as(u5, @intCast(self.height - @as(i32, @intCast(x)) - 1)))
);
} else {
return res;
}
}
};
// binary search to find index in sorted array
fn find(arr: []const i32, x: i32) usize {
var left: usize = 0;
var right: usize = arr.len;
while (left < right) {
const mid = (left + right) / 2;
if (arr[mid] < x) {
left = mid + 1;
} else {
right = mid;
}
}
return left;
}
pub fn main() !void {
var gpa = std.heap.GeneralPurposeAllocator(.{}){};
defer _ = gpa.deinit();
const allocator = gpa.allocator();
const n: usize = 5;
const a = [_]i32{ 3374, 956, 2114, 3415, 3437 };
var input = try allocator.alloc(i32, n);
defer allocator.free(input);
@memcpy(input, &a);
const backup = try allocator.alloc(i32, n);
defer allocator.free(backup);
@memcpy(backup, &a);
// Sort and deduplicate the array
const sorted_a = try allocator.alloc(i32, n);
defer allocator.free(sorted_a);
@memcpy(sorted_a, &a);
std.sort.insertion(i32, sorted_a, {}, std.sort.asc(i32));
// Deduplicate
var unique_a = std.ArrayList(i32).init(allocator);
defer unique_a.deinit();
for (sorted_a, 0..) |val, i| {
if (i == 0 or val != sorted_a[i-1]) {
try unique_a.append(val);
}
}
// Map original values to their indices in the unique array
for (0..n) |i| {
input[i] = @as(i32, @intCast(find(unique_a.items, backup[i])));
}
const lrk_vector = [_][3]i32{
[_]i32{ 2, 2, 1 },
[_]i32{ 3, 4, 1 },
[_]i32{ 4, 5, 1 },
[_]i32{ 1, 2, 2 },
[_]i32{ 4, 4, 1 },
};
var wm = try WaveletMatrix.new(allocator, input, null);
defer wm.deinit();
const stdout = std.io.getStdOut().writer();
for (lrk_vector) |lrk| {
const l = lrk[0] - 1; // Convert to 0-indexed
const r = lrk[1];
const k = lrk[2];
try stdout.print("{d}\n", .{unique_a.items[@as(usize, @intCast(wm.quantile(k - 1, @as(usize, @intCast(l)), @as(usize, @intCast(r)))))]});
}
}
- Output:
956 2114 3415 3374 3415