\begin{algorithm}[!bth] 
\caption{Rescale to budget}
\label{alg:rescale}

\KwInput{Predicted probabilities $\vec{p} \in \mathbb{R}^{n_a}$, Sampling budget \texttt{budget}} \KwOutput{Rescaled probabilities $\vec{p}_{\text{rescaled}}$}


$s \gets \texttt{budget} / n_a$ \tcp*{Calculate target sparsity based on budget}

$\bar{p} \gets \mathbb{E}(\vec{p})$ \tcp*{Compute mean probability}

$r \gets s / \bar{p}$ \tcp*{Scaling factor for each sample}

$\beta \gets (1 - s) / (1 - \bar{p})$ \tcp*{Adjustment factor for non-scaled elements}

\If{$r \leq 1$}
    {$\vec{p}_{\text{rescaled}} \gets \vec{p} \cdot r$  \tcp*{Scale down to fit target sparsity}}
\Else 
    {$\vec{p}_{\text{rescaled}} \gets \boldsymbol{1} - (\boldsymbol{1} - \vec{p}) \times \beta$
    \tcp*{Adjust larger values while preserving distribution}}


\textbf{return} {$\vec{p}_{\text{rescaled}}$ such that $\mathbb{E}(\vec{p}_{\text{rescaled}}) = s$} \end{algorithm}


