Weighted sampling library

Ziwen Wang

February 23, 2026

In this writing we will study practically efficient method for weighted sampling.

1 Theoretical discussion

Definition 1 (WSS). Let $S$ be an input set of $n$ element. For each $a_{i} \in S$ we define $w (a_{i}) \in ℝ_{+}$ to be the weight of element $a_{i}$ . Given an input $t \in ℤ_{+}$ the WSS returns $t$ -independent samples from $S$ , each element $a_{i}$ is sampled with respect to $\frac{w (a_{i})}{w (S)}$ , where $w (S) : = \sum_{a_{i} \in S} w (a_{i})$ .

It is not hard to see by using a standard balanced binary search tree method one can solve such a problem in $O (tlog (n))$ time, with preprocess time $O (nlog (n))$ and space compelxity $O (n)$ . The advantage of using a tree is that updating the weight becomes trivial in $O (\log (n))$ time. On the other hand one can use $n$ bucket’s to achieve $O (t)$ time for sampling $O (t)$ items, but this method does not support weight update, if a weight of an element changed then the whole data structure have to be built in $O (n)$ time.

The question is then can we combine the advantages together? i.e. Does there exists an data structure that solves the above problem in $O (t)$ time with $O (1)$ update, $O (n)$ space complexity and, $O (n)$ preprocessing time? The answer is yes, theoretically.

Theorem 2. There exists a data structures that solves the WSS problem in $O (t)$ time for $t$ samples. The required preprocessing time for $n$ element is $O (n)$ with space complexity $O (n)$ . The update requires $O (1)$ time.

The above method is not good in practice, the reaons is because they have designed nested levels of bucketing and as a result it has a large hidden constant factor makeing it impractical.

1.1 Bucket sampling

Suppose we have a set of buckets, each bucket $B_{i}$ contains element with weight ranging in $[2^{i}, 2^{i + 1}]$ . One subproblem we want to solve is given a bucket $B_{i}$ how would we sample an element $a \in B_{i}$ with probability $\frac{w (a)}{\sum_{b \in B_{i}} w (b)}$ in constant time? The idea is the following procedure,

1.: $x \leftarrow rand (0, 2^{i + 1})$ , sample a point from $B_{i}$ uniformly.
2.: if $x \leq w (a)$ then return $a$ else repeat.

Lemma 3. The above procedure ends in $O (1)$ expected time and samples $a$ with probability $\frac{w (a)}{\sum_{b \in B_{i}} w (b)}$ .

Proof. Since $w (a) \geq 2^{i}$ and $x \leftarrow ran (0, 2^{i + 1})$ then by expectation we know that we should end within $2$ iterations in expectation.

Suppose we are at iteration $t$ , then the probability that we accept $a$ is the product of the probability that $x \leq w (a)$ and we sampled $a$ uniformly from $B_{i}$ which is exactly $\frac{w (a)}{2^{i + 1} | B_{i} |}$ . The probabilty that we return in iteration $t$ is $\sum_{b \in B_{i}} \Pr (b is returned) = \frac{\sum_{b \in B_{i}}}{2^{i + 1} | B_{i} |}$ . Then letting $p = \frac{\sum_{b \in B_{i}}}{2^{i + 1} | B_{i} |}$ and $q = \frac{w (a)}{2^{i + 1} | B_{i} |}$ we see that the probability we sampled $a$ is:

\sum_{t = 1}^{\infty} {(1 - p)}^{t - 1} q = q \sum_{t = 1}^{\infty} = q ∕ p = \frac{w (a)}{\sum_{b \in B_{i}} w (b)},

as required. □

So now the challenge is actually how do we sampled the buckets i.e. each bucket is sampled with probabilty $\frac{w (B_{i})}{w (S)}$ . The idea is to actually sample $x \leftarrow rand (0, w (S))$ and then examine from the high weight buckets to the lower weight bucket one by one, in the worst case we may examine $n$ buckets, but this happens rarely as summarized by the following lemma.

Lemma 4. Let $r$ be the largest weight non-empty bucket $B_{r}$ , then the sum of weights falling into buckets smaller than $r - 2 ⌈ \log (n) ⌉$ is at most $\frac{w (S)}{n + 1}$ .

By the above lemma we know with probability $\frac{n - 1}{n + 1}$ we are going to incur a cost of $O (\log (n))$ by examining buckets in $[r - 2 ⌊ \log (n) ⌋, r]$ and with probability $\frac{1}{n + 1}$ we are going to incur a cost of $O (n)$ to examine buckets $< r - 2 ⌊ \log (n) ⌋$ , so the amortized time is $O (\log (n))$ . A naive algorithm would be,

1.: $x \leftarrow rand (0, w (S))$ do a linear scan to find a bucket $B_{i}$ such that $\sum_{j < i} w (B_{j}) < x \leq \sum_{j \leq i} w (B_{j})$ .
2.: Use previously mentioned procedure to sample an element $a$ from $B_{i}$ in constant time.

It is not hard to see that the above procedure samples $t$ element in $O (tlog (n))$ amortized time and the reason for this is because we need to incur $O (\log (n))$ cost to first find the desired bucket, can we lower this down to $O (1)$ amortized time? The answer is yes, first let the cut off point be $C : = \sum_{i < r - 2 ⌊ \log (n) ⌋} w (B_{i})$ then if we sampled $x \leftarrow rand (0, w (S))$ such that $x \geq C$ we are in the good case and otherwise we are in the bad case. For the bad case this only happens $\leq \frac{1}{n}$ time and incurs a cost of $O (1)$ and for the good case we can build an Alias data structure on non-empty buckets from $[B_{r - 2 ⌊ \log (n) ⌋}, . . ., B_{r}]$ with weight $[w (B_{r - 2 ⌊ \log (n) ⌋}), . . ., w (B_{r})]$ which costs $O (\log (n))$ space and time to preprocess, and once it is preprocessed it only requires $O (1)$ time to sample an bucket. Thus we have the following procedure,

1.: Build an Alias data structure on nonempty buckets ranging from $[r - 2 ⌊ \log (n) ⌋, r]$ , and let $C$ be the cutoff point as defined above.
2.: $x \leftarrow rand (0, w (S))$ . If $x < C$ then we are in the bad case do a linear scan to find the bucket else sample a bucket from the Alias data strcture. Call this bucket $B_{i}$ .
3.: Use previously mentioned procedure to sample an element $a$ from $B_{i}$ in constant time.
4.: Repeat until we have $t$ samples.

The above procedure samples $t$ elements in $O (\log (n) + t)$ time. For updating the weight of an element $a \in S$ we first need remove $a$ and insert $a$ into the appropiate bucket which can be done in $O (1)$ time by book keeping where $a$ is stored at. The value $r$ , largest non-empty bucket $B_{r}$ , can also be book keeped easily if we only allow updates i.e. no insert or delete. Once we allow insert or delete it is not so clear how we can book keep $r$ in $O (1)$ time but we can indeed book keep in $O (\log (n))$ time by balanced binary search tree, another approach is to recalculate $r$ in $O (\log (N))$ during the sampling procedure so update (including insert and delete) stays in $O (1)$ time which is summarized by the following lemma.

Lemma 5. Let $w (S)$ be the sum of all weights in $S$ , let $r^{'} = \log (w (S))$ then the largest non-empty bucket $B_{r}$ has $r \in [r^{'} - ⌈ \log (n) ⌉, r^{'}]$ .

Theorem 6. [ZJW23] There exists a data structures that solves the WSS problem in $O (t + \log (n))$ amortized time for $t$ samples. The required preprocessing time for $n$ element is $O (n)$ with space complexity $O (n)$ . The update requires $O (1)$ time. NOTE: The hidden constants are small which makes this data structure practical.

2 Implementation

Note that the author of the paper has its own implementation but there are several cavaets, the first one being we aim for a much simpler and clean implementation for maintiability and usability, for detail please see WSS class below.

1#ifndef WSS_H 
2#define WSS_H 
3 
4#include <concepts> 
5#include <vector> 
6#include <random> 
7#include <assert.h> 
8 
9#define EPS 1e-6 
10 
11template <typename U> 
12concept Numeric = std::integral<U> || std::floating_point<U>; 
13 
14template <typename C, typename U> 
15concept ContainerOf = requires(C c, std::size_t i) { 
16    { c.begin() }; 
17    { c.end() }; 
18    { c.size() } -> std::convertible_to<std::size_t>; 
19    { c[i] } -> std::convertible_to<U>; 
20    { c.data() } -> std::convertible_to<U*>; 
21}; 
22 
23template <Numeric T> 
24T rand(T l, T r) { 
25    static thread_local std::mt19937 gen{std::random_device{}()}; 
26    if constexpr (std::integral<T>) { 
27        std::uniform_int_distribution<T> dist(l, r); 
28        return dist(gen); 
29    } else { 
30        std::uniform_real_distribution<T> dist(l, r); 
31        return dist(gen); 
32    } 
33} 
34 
35template std::size_t rand<std::size_t>(std::size_t, std::size_t); 
36template int rand<int>(int, int); 
37template double rand<double>(double, double); 
38 
39template <Numeric T, ContainerOf<T> C =  std::vector<T>> 
40class WSS { 
41public: 
42    WSS(const C &weight) : weight_(weight), n_(weight.size()) { 
43        assert("weight␣cannot␣have␣size␣0" && n_); 
44    } 
45    virtual ~WSS() = default; 
46 
47    virtual std::vector<std::size_t> sample(std::size_t m) = 0; 
48    virtual std::size_t sample() = 0; 
49    virtual bool update(std::size_t idx, T weight) = 0; 
50protected: 
51    C weight_; 
52    std::size_t n_; 
53}; 
54 
55#endif //WSS_H

1#ifndef ALIAS 
2#define ALIAS 
3 
4#include <assert.h> 
5#include <vector> 
6#include <optional> 
7#include <utility> 
8#include <stdexcept> 
9#include <random> 
10 
11#include "wss.h" 
12 
13template<typename T> 
14struct Node { 
15    T val; 
16    std::size_t id; 
17    Node *next = nullptr; 
18    Node *prev = nullptr; 
19}; 
20 
21template<typename T> 
22struct OptionalPair { 
23    // First is idx and second is weight. 
24    std::pair<T, T> p1; 
25    std::optional<std::pair<T, T>> p2; 
26}; 
27 
28template<Numeric T, ContainerOf<T> C = std::vector<T>> 
29class Alias : public WSS<T, C> { 
30public: 
31    Alias(const C &weight); 
32 
33    std::vector<std::size_t> sample(std::size_t m) override; 
34    std::size_t sample() override; 
35    bool update(std::size_t idx, T weight) override { 
36        throw std::logic_error("update()␣not␣supported␣for␣Alias"); 
37    } 
38private: 
39    double average_; 
40    std::vector<OptionalPair<T>> bins_; 
41}; 
42 
43template<Numeric T, ContainerOf<T> C> 
44Alias<T, C>::Alias(const C &weight) : WSS<T, C>(weight) { 
45    average_ = 0; 
46    for (auto w : this->weight_) { 
47        average_ += w; 
48    } 
49    average_ /= this->n_; 
50 
51    Node<T> *Less = nullptr; 
52    Node<T> *More = nullptr; 
53    std::size_t idx = 0; 
54    for (auto w : this->weight_) { 
55        if (w <= average_) { 
56          Less = new Node{w, idx, Less}; 
57        } else { 
58          More = new Node{w, idx, More}; 
59        } 
60        ++idx; 
61    } 
62 
63    bins_.reserve(this->n_); 
64    auto nextAndDelete = [](Node<T> *node) { 
65        delete node; 
66    }; 
67    while(Less != nullptr) { 
68        Node<T> *ToDelete = Less; 
69        if (Less->val < average_) { 
70            T lessVal = Less->val; 
71            bins_.push_back(OptionalPair<T>{std::make_pair(Less->id, lessVal), More != nullptr ? std::make_optional(std::make_pair(More->id, average_ - lessVal)) : std::nullopt}); 
72            Less = Less->next; 
73            if (More) { 
74                More->val -= average_ - lessVal; 
75                if (More->val <= average_ + EPS) { 
76                    Node<T> *tmp = More; 
77                    More = More->next; 
78                    tmp->next = Less; 
79                    Less = tmp; 
80                } 
81            } 
82        } else { 
83            bins_.push_back(OptionalPair<T>{std::make_pair(Less->id, Less->val), std::nullopt}); 
84            Less = Less->next; 
85        } 
86        nextAndDelete(ToDelete); 
87    } 
88 
89    Node<T> *tmp = More; 
90    while (tmp != nullptr) { 
91        tmp = tmp->next; 
92    } 
93    assert(More == nullptr && Less == nullptr); 
94} 
95 
96 
97template<Numeric T, ContainerOf<T> C> 
98std::vector<std::size_t> Alias<T, C>::sample(std::size_t m) { 
99    std::vector<std::size_t> res(m, 0); 
100    for (std::size_t i = 0; i < m; ++i) { 
101        res[i] = sample(); 
102    } 
103    return res; 
104} 
105 
106template<Numeric T, ContainerOf<T> C> 
107std::size_t Alias<T, C>::sample() { 
108    std::uniform_int_distribution<std::size_t> binDist(0, this->n_ - 1); 
109    OptionalPair<T> p = bins_[rand(static_cast<std::size_t>(0), this->n_ - 1)]; 
110 
111    if (!p.p2.has_value()) { 
112        return p.p1.first; 
113    } 
114 
115    T r; 
116    r = rand(static_cast<T>(0), static_cast<T>(average_)); 
117 
118    if (r <= p.p1.second) return p.p1.first; 
119    else return p.p2->first; 
120} 
121 
122#endif //ALIAS

1#ifndef BUS 
2#define BUS 
3 
4#include <map> 
5#include <optional> 
6#include <assert.h> 
7 
8#include "wss.h" 
9#include "alias.h" 
10 
11#define max(a, b) (a > b) ? a : b 
12 
13template<Numeric T, ContainerOf<T> C = std::vector<T>> 
14class Bus : public WSS<T, C> { 
15public: 
16    Bus(const C &weight); 
17 
18    std::vector<std::size_t> sample(std::size_t m) override; 
19    std::size_t sample() override; 
20    bool update(std::size_t idx, T weight) override; 
21    std::size_t getBin(std::size_t ele) { 
22        return static_cast<size_t>(floor(std::log2(this->weight_[ele]))); 
23    } 
24 
25    int getMaxNonEmptyBucket(); 
26 
27    struct BucketType { 
28        std::vector<std::size_t> ele; 
29        T totalW; 
30    }; 
31private: 
32    std::vector<std::size_t> elementToIdx_; 
33    std::map<std::size_t, BucketType> bins_; 
34    T totalW_; 
35    std::optional<int> maxPower_; 
36}; 
37 
38template<Numeric T, ContainerOf<T> C> 
39Bus<T, C>::Bus(const C& weight) : WSS<T, C>(weight) { 
40    elementToIdx_.assign(weight.size(), 0); 
41    for (std::size_t i = 0; i < weight.size(); ++i) { 
42        T w = weight[i]; 
43        const int power = floor(std::log2(w)); 
44        if (!bins_.contains(power)) { 
45            bins_[power] = BucketType(std::vector<std::size_t>{}, 0); 
46        } 
47 
48        totalW_ += w; 
49        bins_[power].totalW += w; 
50        elementToIdx_[i] = bins_[power].ele.size(); 
51        bins_[power].ele.push_back(i); 
52        maxPower_ = max(power, maxPower_); 
53    } 
54} 
55 
56template<Numeric T, ContainerOf<T> C> 
57int Bus<T, C>::getMaxNonEmptyBucket() { 
58    if (maxPower_) return maxPower_.value(); 
59 
60    int R = floor(std::log2(totalW_)); 
61    int L = R - floor(std::log2(this->n_)); 
62    for (int idx = R; idx >= L; --idx) { 
63        if (bins_.count(idx)) { maxPower_ = idx; break; } 
64    } 
65    assert("Bug:␣maxPower_␣is␣empty" && maxPower_); 
66    return maxPower_.value(); 
67} 
68 
69template<Numeric T, ContainerOf<T> C> 
70std::size_t Bus<T, C>::sample() { 
71    return sample(1)[0]; 
72} 
73 
74template<Numeric T, ContainerOf<T> C> 
75std::vector<std::size_t> Bus<T, C>::sample(std::size_t m) { 
76    std::vector<size_t> res; res.reserve(m); 
77    int r = getMaxNonEmptyBucket(); 
78 
79    T cutOff = 0; 
80    std::vector<T> aliasWeight; 
81    std::vector<std::size_t> idxToPower; 
82 
83    for (int i = r - 2*floor(std::log2(this->n_)); i <= r; ++i) { 
84        if (!bins_.count(i)) continue; 
85        cutOff += bins_[i].totalW; 
86        aliasWeight.push_back(bins_[i].totalW); 
87        idxToPower.push_back(i); 
88    } 
89    cutOff = totalW_ - cutOff; 
90    Alias<T> alias(aliasWeight); 
91 
92    while(m--) { 
93        T wI = rand(0.0, totalW_), preSumWI = cutOff; 
94        bool found = false; 
95        int idx; 
96 
97        // With high probability that we will find idx >= r - (r ? floor(2 * std::log2(r)) : 0). 
98        // Build the Alias method data structure for for powers with [r - 2logn, ..., r]. 
99        // For alias method: 
100        //       1. Find the cutoff weight between buckets of power r-2log(r) - 1 and r - 2log(r). 
101        //       2. Sample a weight from [0, totalW_]. 
102        //       3. If the sampled weight is in good case i.e. greater than cut off point then 
103        //          we use alias method on buckets of power [r-2log n, ..., r] which computes in O(1) time. 
104        //       4. Otherwise use brute force on buckets of power [lowest, ..., r - 2logn] which happens 
105        //          with probability O(1/n) with time cost O(n). 
106        // Using the above described sampling technique yields an amortized time of O(log(n) + m) for m samples. 
107        if (wI >= cutOff) { 
108            idx = idxToPower[alias.sample()]; 
109        } else { 
110            idx = r - 2*floor(std::log2(this->n_)) - 1; 
111            for (; ; --idx) { 
112                if (!bins_.count(idx)) continue; 
113                if (wI >= preSumWI - bins_[idx].totalW) { 
114                    found = 1; break; 
115                } 
116                preSumWI = preSumWI - bins_[idx].totalW; 
117            } 
118        } 
119 
120        const std::vector<std::size_t> &ele = bins_[idx].ele; 
121        while (true) { 
122            T x = rand(static_cast<T>(0), static_cast<T>(1ULL << (idx+1))); 
123            std::size_t eleIdx = ele[rand(static_cast<std::size_t>(0), ele.size() - 1)]; 
124            if (x <= this->weight_[eleIdx]) { res.push_back(eleIdx); break; } 
125        } 
126    } 
127 
128    return res; 
129} 
130 
131template<Numeric T, ContainerOf<T> C> 
132bool Bus<T, C>::update(std::size_t ele, T weight) { 
133    int binIdx = getBin(ele); 
134    int insertIdx = floor(std::log2(weight)); 
135 
136    bins_[binIdx].ele[elementToIdx_[ele]] = bins_[binIdx].ele.back(); 
137    elementToIdx_[bins_[binIdx].ele.back()] = elementToIdx_[ele]; 
138    bins_[binIdx].ele.pop_back(); 
139    bins_[binIdx].totalW -= this->weight_[ele]; 
140 
141    if (bins_[binIdx].ele.size() == 0) bins_.erase(binIdx); 
142 
143    elementToIdx_[ele] = bins_[insertIdx].ele.size(); 
144    bins_[insertIdx].ele.push_back(ele); 
145    bins_[insertIdx].totalW += weight; 
146 
147    totalW_ += weight - this->weight_[ele]; 
148    this->weight_[ele] = weight; 
149 
150 
151    if (binIdx <= insertIdx) { 
152        maxPower_ = max(insertIdx, maxPower_); 
153    } else { 
154        maxPower_.reset(); 
155    } 
156 
157    return true; 
158} 
159 
160#endif //BUS

References

[ZJW23] Fangyuan Zhang, Mengxu Jiang, and Sibo Wang. Efficient dynamic weighted set sampling and its extension. 17(1), 2023.