commons-rng-sampling/src/main/java/org/apache/commons/rng/sampling/distribution/FastLoadedDiceRollerDiscreteSampler.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.commons.rng.sampling.distribution;

import java.math.BigInteger;
import java.util.Arrays;
import org.apache.commons.rng.UniformRandomProvider;

/**
 * Distribution sampler that uses the Fast Loaded Dice Roller (FLDR). It can be used to
 * sample from {@code n} values each with an associated relative weight. If all unique items
 * are assigned the same weight it is more efficient to use the {@link DiscreteUniformSampler}.
 *
 * <p>Given a list {@code L} of {@code n} positive numbers,
 * where {@code L[i]} represents the relative weight of the {@code i}th side, FLDR returns
 * integer {@code i} with relative probability {@code L[i]}.
 *
 * <p>FLDR produces <em>exact</em> samples from the specified probability distribution.
 * <ul>
 *   <li>For integer weights, the probability of returning {@code i} is precisely equal to the
 *   rational number {@code L[i] / m}, where {@code m} is the sum of {@code L}.
 *   <li>For floating-points weights, each weight {@code L[i]} is converted to the
 *   corresponding rational number {@code p[i] / q[i]} where {@code p[i]} is a positive integer and
 *   {@code q[i]} is a power of 2. The rational weights are then normalized (exactly) to sum to unity.
 * </ul>
 *
 * <p>Note that if <em>exact</em> samples are not required then an alternative sampler that
 * ignores very small relative weights may have improved sampling performance.
 *
 * <p>This implementation is based on the algorithm in:
 *
 * <blockquote>
 *  Feras A. Saad, Cameron E. Freer, Martin C. Rinard, and Vikash K. Mansinghka.
 *  The Fast Loaded Dice Roller: A Near-Optimal Exact Sampler for Discrete Probability
 *  Distributions. In AISTATS 2020: Proceedings of the 23rd International Conference on
 *  Artificial Intelligence and Statistics, Proceedings of Machine Learning Research 108,
 *  Palermo, Sicily, Italy, 2020.
 * </blockquote>
 *
 * <p>Sampling uses {@link UniformRandomProvider#nextInt()} as the source of random bits.
 *
 * @see <a href="https://arxiv.org/abs/2003.03830">Saad et al (2020)
 * Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics,
 * PMLR 108:1036-1046.</a>
 * @since 1.5
 */
public abstract class FastLoadedDiceRollerDiscreteSampler
    implements SharedStateDiscreteSampler {
    /**
     * The maximum size of an array.
     *
     * <p>This value is taken from the limit in Open JDK 8 {@code java.util.ArrayList}.
     * It allows VMs to reserve some header words in an array.
     */
    private static final int MAX_ARRAY_SIZE = Integer.MAX_VALUE - 8;
    /** The maximum biased exponent for a finite double.
     * This is offset by 1023 from {@code Math.getExponent(Double.MAX_VALUE)}. */
    private static final int MAX_BIASED_EXPONENT = 2046;
    /** Size of the mantissa of a double. Equal to 52 bits. */
    private static final int MANTISSA_SIZE = 52;
    /** Mask to extract the 52-bit mantissa from a long representation of a double. */
    private static final long MANTISSA_MASK = 0x000f_ffff_ffff_ffffL;
    /** BigInteger representation of {@link Long#MAX_VALUE}. */
    private static final BigInteger MAX_LONG = BigInteger.valueOf(Long.MAX_VALUE);
    /** The maximum offset that will avoid loss of bits for a left shift of a 53-bit value.
     * The value will remain positive for any shift {@code <=} this value. */
    private static final int MAX_OFFSET = 10;
    /** Initial value for no leaf node label. */
    private static final int NO_LABEL = Integer.MAX_VALUE;
    /** Name of the sampler. */
    private static final String SAMPLER_NAME = "Fast Loaded Dice Roller";

    /**
     * Class to handle the edge case of observations in only one category.
     */
    private static class FixedValueDiscreteSampler extends FastLoadedDiceRollerDiscreteSampler {
        /** The sample value. */
        private final int sampleValue;

        /**
         * @param sampleValue Sample value.
         */
        FixedValueDiscreteSampler(int sampleValue) {
            this.sampleValue = sampleValue;
        }

        @Override
        public int sample() {
            return sampleValue;
        }

        @Override
        public FastLoadedDiceRollerDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
            return this;
        }

        @Override
        public String toString() {
            return SAMPLER_NAME;
        }
    }

    /**
     * Class to implement the FLDR sample algorithm.
     */
    private static class FLDRSampler extends FastLoadedDiceRollerDiscreteSampler {
        /** Empty boolean source. This is the location of the sign-bit after 31 right shifts on
         * the boolean source. */
        private static final int EMPTY_BOOL_SOURCE = 1;

        /** Underlying source of randomness. */
        private final UniformRandomProvider rng;
        /** Number of categories. */
        private final int n;
        /** Number of levels in the discrete distribution generating (DDG) tree.
         * Equal to {@code ceil(log2(m))} where {@code m} is the sum of observations. */
        private final int k;
        /** Number of leaf nodes at each level. */
        private final int[] h;
        /** Stores the leaf node labels in increasing order. Named {@code H} in the FLDR paper. */
        private final int[] lH;

        /**
         * Provides a bit source for booleans.
         *
         * <p>A cached value from a call to {@link UniformRandomProvider#nextInt()}.
         *
         * <p>Only stores 31-bits when full as 1 bit has already been consumed.
         * The sign bit is a flag that shifts down so the source eventually equals 1
         * when all bits are consumed and will trigger a refill.
         */
        private int booleanSource = EMPTY_BOOL_SOURCE;

        /**
         * Creates a sampler.
         *
         * <p>The input parameters are not validated and must be correctly computed tables.
         *
         * @param rng Generator of uniformly distributed random numbers.
         * @param n Number of categories
         * @param k Number of levels in the discrete distribution generating (DDG) tree.
         * Equal to {@code ceil(log2(m))} where {@code m} is the sum of observations.
         * @param h Number of leaf nodes at each level.
         * @param lH Stores the leaf node labels in increasing order.
         */
        FLDRSampler(UniformRandomProvider rng,
                    int n,
                    int k,
                    int[] h,
                    int[] lH) {
            this.rng = rng;
            this.n = n;
            this.k = k;
            // Deliberate direct storage of input arrays
            this.h = h;
            this.lH = lH;
        }

        /**
         * Creates a copy with a new source of randomness.
         *
         * @param rng Generator of uniformly distributed random numbers.
         * @param source Source to copy.
         */
        private FLDRSampler(UniformRandomProvider rng,
                            FLDRSampler source) {
            this.rng = rng;
            this.n = source.n;
            this.k = source.k;
            this.h = source.h;
            this.lH = source.lH;
        }

        /** {@inheritDoc} */
        @Override
        public int sample() {
            // ALGORITHM 5: SAMPLE
            int c = 0;
            int d = 0;
            for (;;) {
                // b = flip()
                // d = 2 * d + (1 - b)
                d = (d << 1) + flip();
                if (d < h[c]) {
                    // z = H[d][c]
                    final int z = lH[d * k + c];
                    // assert z != NO_LABEL
                    if (z < n) {
                        return z;
                    }
                    d = 0;
                    c = 0;
                } else {
                    d = d - h[c];
                    c++;
                }
            }
        }

        /**
         * Provides a source of boolean bits.
         *
         * <p>Note: This replicates the boolean cache functionality of
         * {@code o.a.c.rng.core.source32.IntProvider}. The method has been simplified to return
         * an {@code int} value rather than a {@code boolean}.
         *
         * @return the bit (0 or 1)
         */
        private int flip() {
            int bits = booleanSource;
            if (bits == 1) {
                // Refill
                bits = rng.nextInt();
                // Store a refill flag in the sign bit and the unused 31 bits, return lowest bit
                booleanSource = Integer.MIN_VALUE | (bits >>> 1);
                return bits & 0x1;
            }
            // Shift down eventually triggering refill, return current lowest bit
            booleanSource = bits >>> 1;
            return bits & 0x1;
        }

        /** {@inheritDoc} */
        @Override
        public String toString() {
            return SAMPLER_NAME + " [" + rng.toString() + "]";
        }

        /** {@inheritDoc} */
        @Override
        public FastLoadedDiceRollerDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
            return new FLDRSampler(rng, this);
        }
    }

    /** Package-private constructor. */
    FastLoadedDiceRollerDiscreteSampler() {
        // Intentionally empty
    }

    /** {@inheritDoc} */
    // Redeclare the signature to return a FastLoadedDiceRollerSampler not a SharedStateLongSampler
    @Override
    public abstract FastLoadedDiceRollerDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng);

    /**
     * Creates a sampler.
     *
     * <p>Note: The discrete distribution generating (DDG) tree requires {@code (n + 1) * k} entries
     * where {@code n} is the number of categories, {@code k == ceil(log2(m))} and {@code m}
     * is the sum of the observed frequencies. An exception is raised if this cannot be allocated
     * as a single array.
     *
     * <p>For reference the sum is limited to {@link Long#MAX_VALUE} and the value {@code k} to 63.
     * The number of categories is limited to approximately {@code ((2^31 - 1) / k) = 34,087,042}
     * when the sum of frequencies is large enough to create k=63.
     *
     * @param rng Generator of uniformly distributed random numbers.
     * @param frequencies Observed frequencies of the discrete distribution.
     * @return the sampler
     * @throws IllegalArgumentException if {@code frequencies} is null or empty, a
     * frequency is negative, the sum of all frequencies is either zero or
     * above {@link Long#MAX_VALUE}, or the size of the discrete distribution generating tree
     * is too large.
     */
    public static FastLoadedDiceRollerDiscreteSampler of(UniformRandomProvider rng,
                                                         long[] frequencies) {
        final long m = sum(frequencies);

        // Obtain indices of non-zero frequencies
        final int[] indices = indicesOfNonZero(frequencies);

        // Edge case for 1 non-zero weight. This also handles edge case for 1 observation
        // (as log2(m) == 0 will break the computation of the DDG tree).
        if (indices.length == 1) {
            return new FixedValueDiscreteSampler(indexOfNonZero(frequencies));
        }

        return createSampler(rng, frequencies, indices, m);
    }

    /**
     * Creates a sampler.
     *
     * <p>Weights are converted to rational numbers {@code p / q} where {@code q} is a power of 2.
     * The numerators {@code p} are scaled to use a common denominator before summing.
     *
     * <p>All weights are used to create the sampler. Weights with a small magnitude relative
     * to the largest weight can be excluded using the constructor method with the
     * relative magnitude parameter {@code alpha} (see {@link #of(UniformRandomProvider, double[], int)}).
     *
     * @param rng Generator of uniformly distributed random numbers.
     * @param weights Weights of the discrete distribution.
     * @return the sampler
     * @throws IllegalArgumentException if {@code weights} is null or empty, a
     * weight is negative, infinite or {@code NaN}, the sum of all weights is zero, or the size
     * of the discrete distribution generating tree is too large.
     * @see #of(UniformRandomProvider, double[], int)
     */
    public static FastLoadedDiceRollerDiscreteSampler of(UniformRandomProvider rng,
                                                         double[] weights) {
        return of(rng, weights, 0);
    }

    /**
     * Creates a sampler.
     *
     * <p>Weights are converted to rational numbers {@code p / q} where {@code q} is
     * a power of 2. The numerators {@code p} are scaled to use a common
     * denominator before summing.
     *
     * <p>Note: The discrete distribution generating (DDG) tree requires
     * {@code (n + 1) * k} entries where {@code n} is the number of categories,
     * {@code k == ceil(log2(m))} and {@code m} is the sum of the weight numerators
     * {@code q}. An exception is raised if this cannot be allocated as a single
     * array.
     *
     * <p>For reference the value {@code k} is equal to or greater than the ratio of
     * the largest to the smallest weight expressed as a power of 2. For
     * {@code Double.MAX_VALUE / Double.MIN_VALUE} this is ~2098. The value
     * {@code k} increases with the sum of the weight numerators. A number of
     * weights in excess of 1,000,000 with values equal to {@link Double#MAX_VALUE}
     * would be required to raise an exception when the minimum weight is
     * {@link Double#MIN_VALUE}.
     *
     * <p>Weights with a small magnitude relative to the largest weight can be
     * excluded using the relative magnitude parameter {@code alpha}. This will set
     * any weight to zero if the magnitude is approximately 2<sup>alpha</sup>
     * <em>smaller</em> than the largest weight. This comparison is made using only
     * the exponent of the input weights. The {@code alpha} parameter is ignored if
     * not above zero. Note that a small {@code alpha} parameter will exclude more
     * weights than a large {@code alpha} parameter.
     *
     * <p>The alpha parameter can be used to exclude categories that
     * have a very low probability of occurrence and will improve the construction
     * performance of the sampler. The effect on sampling performance depends on
     * the relative weights of the excluded categories; typically a high {@code alpha}
     * is used to exclude categories that would be visited with a very low probability
     * and the sampling performance is unchanged.
     *
     * <p><b>Implementation Note</b>
     *
     * <p>This method creates a sampler with <em>exact</em> samples from the
     * specified probability distribution. It is recommended to use this method:
     * <ul>
     *  <li>if the weights are computed, for example from a probability mass function; or
     *  <li>if the weights sum to an infinite value.
     * </ul>
     *
     * <p>If the weights are computed from empirical observations then it is
     * recommended to use the factory method
     * {@link #of(UniformRandomProvider, long[]) accepting frequencies}. This
     * requires the total number of observations to be representable as a long
     * integer.
     *
     * <p>Note that if all weights are scaled by a power of 2 to be integers, and
     * each integer can be represented as a positive 64-bit long value, then the
     * sampler created using this method will match the output from a sampler
     * created with the scaled weights converted to long values for the factory
     * method {@link #of(UniformRandomProvider, long[]) accepting frequencies}. This
     * assumes the sum of the integer values does not overflow.
     *
     * <p>It should be noted that the conversion of weights to rational numbers has
     * a performance overhead during construction (sampling performance is not
     * affected). This may be avoided by first converting them to integer values
     * that can be summed without overflow. For example by scaling values by
     * {@code 2^62 / sum} and converting to long by casting or rounding.
     *
     * <p>This approach may increase the efficiency of construction. The resulting
     * sampler may no longer produce <em>exact</em> samples from the distribution.
     * In particular any weights with a converted frequency of zero cannot be
     * sampled.
     *
     * @param rng Generator of uniformly distributed random numbers.
     * @param weights Weights of the discrete distribution.
     * @param alpha Alpha parameter.
     * @return the sampler
     * @throws IllegalArgumentException if {@code weights} is null or empty, a
     * weight is negative, infinite or {@code NaN}, the sum of all weights is zero,
     * or the size of the discrete distribution generating tree is too large.
     * @see #of(UniformRandomProvider, long[])
     */
    public static FastLoadedDiceRollerDiscreteSampler of(UniformRandomProvider rng,
                                                         double[] weights,
                                                         int alpha) {
        final int n = checkWeightsNonZeroLength(weights);

        // Convert floating-point double to a relative weight
        // using a shifted integer representation
        final long[] frequencies = new long[n];
        final int[] offsets = new int[n];
        convertToIntegers(weights, frequencies, offsets, alpha);

        // Obtain indices of non-zero weights
        final int[] indices = indicesOfNonZero(frequencies);

        // Edge case for 1 non-zero weight.
        if (indices.length == 1) {
            return new FixedValueDiscreteSampler(indexOfNonZero(frequencies));
        }

        final BigInteger m = sum(frequencies, offsets, indices);

        // Use long arithmetic if possible. This occurs when the weights are similar in magnitude.
        if (m.compareTo(MAX_LONG) <= 0) {
            // Apply the offset
            for (int i = 0; i < n; i++) {
                frequencies[i] <<= offsets[i];
            }
            return createSampler(rng, frequencies, indices, m.longValue());
        }

        return createSampler(rng, frequencies, offsets, indices, m);
    }

    /**
     * Sum the frequencies.
     *
     * @param frequencies Frequencies.
     * @return the sum
     * @throws IllegalArgumentException if {@code frequencies} is null or empty, a
     * frequency is negative, or the sum of all frequencies is either zero or above
     * {@link Long#MAX_VALUE}
     */
    private static long sum(long[] frequencies) {
        // Validate
        if (frequencies == null || frequencies.length == 0) {
            throw new IllegalArgumentException("frequencies must contain at least 1 value");
        }

        // Sum the values.
        // Combine all the sign bits in the observations and the intermediate sum in a flag.
        long m = 0;
        long signFlag = 0;
        for (final long o : frequencies) {
            m += o;
            signFlag |= o | m;
        }

        // Check for a sign-bit.
        if (signFlag < 0) {
            // One or more observations were negative, or the sum overflowed.
            for (final long o : frequencies) {
                if (o < 0) {
                    throw new IllegalArgumentException("frequencies must contain positive values: " + o);
                }
            }
            throw new IllegalArgumentException("Overflow when summing frequencies");
        }
        if (m == 0) {
            throw new IllegalArgumentException("Sum of frequencies is zero");
        }
        return m;
    }

    /**
     * Convert the floating-point weights to relative weights represented as
     * integers {@code value * 2^exponent}. The relative weight as an integer is:
     *
     * <pre>
     * BigInteger.valueOf(value).shiftLeft(exponent)
     * </pre>
     *
     * <p>Note that the weights are created using a common power-of-2 scaling
     * operation so the minimum exponent is zero.
     *
     * <p>A positive {@code alpha} parameter is used to set any weight to zero if
     * the magnitude is approximately 2<sup>alpha</sup> <em>smaller</em> than the
     * largest weight. This comparison is made using only the exponent of the input
     * weights.
     *
     * @param weights Weights of the discrete distribution.
     * @param values Output floating-point mantissas converted to 53-bit integers.
     * @param exponents Output power of 2 exponent.
     * @param alpha Alpha parameter.
     * @throws IllegalArgumentException if a weight is negative, infinite or
     * {@code NaN}, or the sum of all weights is zero.
     */
    private static void convertToIntegers(double[] weights, long[] values, int[] exponents, int alpha) {
        int maxExponent = Integer.MIN_VALUE;
        for (int i = 0; i < weights.length; i++) {
            final double weight = weights[i];
            // Ignore zero.
            // When creating the integer value later using bit shifts the result will remain zero.
            if (weight == 0) {
                continue;
            }
            final long bits = Double.doubleToRawLongBits(weight);

            // For the IEEE 754 format see Double.longBitsToDouble(long).

            // Extract the exponent (with the sign bit)
            int exp = (int) (bits >>> MANTISSA_SIZE);
            // Detect negative, infinite or NaN.
            // Note: Negative values sign bit will cause the exponent to be too high.
            if (exp > MAX_BIASED_EXPONENT) {
                throw new IllegalArgumentException("Invalid weight: " + weight);
            }
            long mantissa;
            if (exp == 0) {
                // Sub-normal number:
                mantissa = (bits & MANTISSA_MASK) << 1;
                // Here we convert to a normalised number by counting the leading zeros
                // to obtain the number of shifts of the most significant bit in
                // the mantissa that is required to get a 1 at position 53 (i.e. as
                // if it were a normal number with assumed leading bit).
                final int shift = Long.numberOfLeadingZeros(mantissa << 11);
                mantissa <<= shift;
                exp -= shift;
            } else {
                // Normal number. Add the implicit leading 1-bit.
                mantissa = (bits & MANTISSA_MASK) | (1L << MANTISSA_SIZE);
            }

            // Here the floating-point number is equal to:
            // mantissa * 2^(exp-1075)

            values[i] = mantissa;
            exponents[i] = exp;
            maxExponent = Math.max(maxExponent, exp);
        }

        // No exponent indicates that all weights are zero
        if (maxExponent == Integer.MIN_VALUE) {
            throw new IllegalArgumentException("Sum of weights is zero");
        }

        filterWeights(values, exponents, alpha, maxExponent);
        scaleWeights(values, exponents);
    }

    /**
     * Filters small weights using the {@code alpha} parameter.
     * A positive {@code alpha} parameter is used to set any weight to zero if
     * the magnitude is approximately 2<sup>alpha</sup> <em>smaller</em> than the
     * largest weight. This comparison is made using only the exponent of the input
     * weights.
     *
     * @param values 53-bit values.
     * @param exponents Power of 2 exponent.
     * @param alpha Alpha parameter.
     * @param maxExponent Maximum exponent.
     */
    private static void filterWeights(long[] values, int[] exponents, int alpha, int maxExponent) {
        if (alpha > 0) {
            // Filter weights. This must be done before the values are shifted so
            // the exponent represents the approximate magnitude of the value.
            for (int i = 0; i < exponents.length; i++) {
                if (maxExponent - exponents[i] > alpha) {
                    values[i] = 0;
                }
            }
        }
    }

    /**
     * Scale the weights represented as integers {@code value * 2^exponent} to use a
     * minimum exponent of zero. The values are scaled to remove any common trailing zeros
     * in their representation. This ultimately reduces the size of the discrete distribution
     * generating (DGG) tree.
     *
     * @param values 53-bit values.
     * @param exponents Power of 2 exponent.
     */
    private static void scaleWeights(long[] values, int[] exponents) {
        // Find the minimum exponent and common trailing zeros.
        int minExponent = Integer.MAX_VALUE;
        for (int i = 0; i < exponents.length; i++) {
            if (values[i] != 0) {
                minExponent = Math.min(minExponent, exponents[i]);
            }
        }
        // Trailing zeros occur when the original weights have a representation with
        // less than 52 binary digits, e.g. {1.5, 0.5, 0.25}.
        int trailingZeros = Long.SIZE;
        for (int i = 0; i < values.length && trailingZeros != 0; i++) {
            trailingZeros = Math.min(trailingZeros, Long.numberOfTrailingZeros(values[i]));
        }
        // Scale by a power of 2 so the minimum exponent is zero.
        for (int i = 0; i < exponents.length; i++) {
            exponents[i] -= minExponent;
        }
        // Remove common trailing zeros.
        if (trailingZeros != 0) {
            for (int i = 0; i < values.length; i++) {
                values[i] >>>= trailingZeros;
            }
        }
    }

    /**
     * Sum the integers at the specified indices.
     * Integers are represented as {@code value * 2^exponent}.
     *
     * @param values 53-bit values.
     * @param exponents Power of 2 exponent.
     * @param indices Indices to sum.
     * @return the sum
     */
    private static BigInteger sum(long[] values, int[] exponents, int[] indices) {
        BigInteger m = BigInteger.ZERO;
        for (final int i : indices) {
            m = m.add(toBigInteger(values[i], exponents[i]));
        }
        return m;
    }

    /**
     * Convert the value and left shift offset to a BigInteger.
     * It is assumed the value is at most 53-bits. This allows optimising the left
     * shift if it is below 11 bits.
     *
     * @param value 53-bit value.
     * @param offset Left shift offset (must be positive).
     * @return the BigInteger
     */
    private static BigInteger toBigInteger(long value, int offset) {
        // Ignore zeros. The sum method uses indices of non-zero values.
        if (offset <= MAX_OFFSET) {
            // Assume (value << offset) <= Long.MAX_VALUE
            return BigInteger.valueOf(value << offset);
        }
        return BigInteger.valueOf(value).shiftLeft(offset);
    }

    /**
     * Creates the sampler.
     *
     * <p>It is assumed the frequencies are all positive and the sum does not
     * overflow.
     *
     * @param rng Generator of uniformly distributed random numbers.
     * @param frequencies Observed frequencies of the discrete distribution.
     * @param indices Indices of non-zero frequencies.
     * @param m Sum of the frequencies.
     * @return the sampler
     */
    private static FastLoadedDiceRollerDiscreteSampler createSampler(UniformRandomProvider rng,
                                                                     long[] frequencies,
                                                                     int[] indices,
                                                                     long m) {
        // ALGORITHM 5: PREPROCESS
        // a == frequencies
        // m = sum(a)
        // h = leaf node count
        // H = leaf node label (lH)

        final int n = frequencies.length;

        // k = ceil(log2(m))
        final int k = 64 - Long.numberOfLeadingZeros(m - 1);
        // r = a(n+1) = 2^k - m
        final long r = (1L << k) - m;

        // Note:
        // A sparse matrix can often be used for H, as most of its entries are empty.
        // This implementation uses a 1D array for efficiency at the cost of memory.
        // This is limited to approximately ((2^31 - 1) / k), e.g. 34087042 when the sum of
        // observations is large enough to create k=63.
        // This could be handled using a 2D array. In practice a number of categories this
        // large is not expected and is currently not supported.
        final int[] h = new int[k];
        final int[] lH = new int[checkArraySize((n + 1L) * k)];
        Arrays.fill(lH, NO_LABEL);

        int d;
        for (int j = 0; j < k; j++) {
            final int shift = (k - 1) - j;
            final long bitMask = 1L << shift;

            d = 0;
            for (final int i : indices) {
                // bool w ← (a[i] >> (k − 1) − j)) & 1
                // h[j] = h[j] + w
                // if w then:
                if ((frequencies[i] & bitMask) != 0) {
                    h[j]++;
                    // H[d][j] = i
                    lH[d * k + j] = i;
                    d++;
                }
            }
            // process a(n+1) without extending the input frequencies array by 1
            if ((r & bitMask) != 0) {
                h[j]++;
                lH[d * k + j] = n;
            }
        }

        return new FLDRSampler(rng, n, k, h, lH);
    }

    /**
     * Creates the sampler. Frequencies are are represented as a 53-bit value with a
     * left-shift offset.
     * <pre>
     * BigInteger.valueOf(value).shiftLeft(offset)
     * </pre>
     *
     * <p>It is assumed the frequencies are all positive.
     *
     * @param rng Generator of uniformly distributed random numbers.
     * @param frequencies Observed frequencies of the discrete distribution.
     * @param offsets Left shift offsets (must be positive).
     * @param indices Indices of non-zero frequencies.
     * @param m Sum of the frequencies.
     * @return the sampler
     */
    private static FastLoadedDiceRollerDiscreteSampler createSampler(UniformRandomProvider rng,
                                                                     long[] frequencies,
                                                                     int[] offsets,
                                                                     int[] indices,
                                                                     BigInteger m) {
        // Repeat the logic from createSampler(...) using extended arithmetic to test the bits

        // ALGORITHM 5: PREPROCESS
        // a == frequencies
        // m = sum(a)
        // h = leaf node count
        // H = leaf node label (lH)

        final int n = frequencies.length;

        // k = ceil(log2(m))
        final int k = m.subtract(BigInteger.ONE).bitLength();
        // r = a(n+1) = 2^k - m
        final BigInteger r = BigInteger.ONE.shiftLeft(k).subtract(m);

        final int[] h = new int[k];
        final int[] lH = new int[checkArraySize((n + 1L) * k)];
        Arrays.fill(lH, NO_LABEL);

        int d;
        for (int j = 0; j < k; j++) {
            final int shift = (k - 1) - j;

            d = 0;
            for (final int i : indices) {
                // bool w ← (a[i] >> (k − 1) − j)) & 1
                // h[j] = h[j] + w
                // if w then:
                if (testBit(frequencies[i], offsets[i], shift)) {
                    h[j]++;
                    // H[d][j] = i
                    lH[d * k + j] = i;
                    d++;
                }
            }
            // process a(n+1) without extending the input frequencies array by 1
            if (r.testBit(shift)) {
                h[j]++;
                lH[d * k + j] = n;
            }
        }

        return new FLDRSampler(rng, n, k, h, lH);
    }

    /**
     * Test the logical bit of the shifted integer representation.
     * The value is assumed to have at most 53-bits of information. The offset
     * is assumed to be positive. This is functionally equivalent to:
     * <pre>
     * BigInteger.valueOf(value).shiftLeft(offset).testBit(n)
     * </pre>
     *
     * @param value 53-bit value.
     * @param offset Left shift offset.
     * @param n Index of bit to test.
     * @return true if the bit is 1
     */
    private static boolean testBit(long value, int offset, int n) {
        if (n < offset) {
            // All logical trailing bits are zero
            return false;
        }
        // Test if outside the 53-bit value (note that the implicit 1 bit
        // has been added to the 52-bit mantissas for 'normal' floating-point numbers).
        final int bit = n - offset;
        return bit <= MANTISSA_SIZE && (value & (1L << bit)) != 0;
    }

    /**
     * Check the weights have a non-zero length.
     *
     * @param weights Weights.
     * @return the length
     */
    private static int checkWeightsNonZeroLength(double[] weights) {
        if (weights == null || weights.length == 0) {
            throw new IllegalArgumentException("weights must contain at least 1 value");
        }
        return weights.length;
    }

    /**
     * Create the indices of non-zero values.
     *
     * @param values Values.
     * @return the indices
     */
    private static int[] indicesOfNonZero(long[] values) {
        int n = 0;
        final int[] indices = new int[values.length];
        for (int i = 0; i < values.length; i++) {
            if (values[i] != 0) {
                indices[n++] = i;
            }
        }
        return Arrays.copyOf(indices, n);
    }

    /**
     * Find the index of the first non-zero frequency.
     *
     * @param frequencies Frequencies.
     * @return the index
     * @throws IllegalStateException if all frequencies are zero.
     */
    static int indexOfNonZero(long[] frequencies) {
        for (int i = 0; i < frequencies.length; i++) {
            if (frequencies[i] != 0) {
                return i;
            }
        }
        throw new IllegalStateException("All frequencies are zero");
    }

    /**
     * Check the size is valid for a 1D array.
     *
     * @param size Size
     * @return the size as an {@code int}
     * @throws IllegalArgumentException if the size is too large for a 1D array.
     */
    static int checkArraySize(long size) {
        if (size > MAX_ARRAY_SIZE) {
            throw new IllegalArgumentException("Unable to allocate array of size: " + size);
        }
        return (int) size;
    }
}