# Efficient computation of the distribution of a weighted sum of finite discrete random variables

I have $N$ discrete random variables $X_i$, with distributions $p_i(X_i)$. $X_i$ can take integer values in a finite range $[a_i,b_i]$, where $a_i,b_i$ are integers. The distribution $p_i(X_i)$ is stored as an array of the values $p_i(a_i)$, $p_i(a_i+1)$, ..., $p_i(b_i)$. Another random variable $Z$ has a distribution $p(z)$ defined as:

$$p(z) \propto \sum_{x_1 = a_1}^{b_1} \dots \sum_{x_N=a_N}^{b_N} \delta \left( cz - \sum_{i=1}^N c_i x_i \right) p_1(x_1) \dots p_N(x_N)$$

where $\delta(x)$ is a Kronecker-delta function, with $\delta(0) = 1$ and $\delta(x) = 0$ for all $x\ne 1$, and $c,c_i$ are known integers. The proportion sign is to allow for a normalization constant.

If $c=c_i=1$ for all $i$, there is a very efficient way to compute the array of values $p(Z)$, using the Fast Fourier Transforms and the Convolution Theorem.

But I am interested in the general case, where $c,c_i$ are arbitrary integers (in my particular scenario they are small integers). Is there a way to adapt the FFT and the convolution theorem for this case too?

I tried making the change of variables $W=c Z$, $Y_i = c_i X_i$, but it gets very messy since you have to account for the fact that since $z,x_i$ have to be integers, $W$ and $Y_i$ will skip many values. I didn't get far down this path. Is this the right way to do it? Is there a better way?

• Your approach is a natural one. Skipping many values should not be a problem for computing the FFTs or their convolutions.
– whuber
Oct 21, 2015 at 15:42
• You can adapt some of the approaches given in stats.stackexchange.com/questions/177163/… Oct 21, 2015 at 17:19
• You can compute this in $O \left( c \times (\max_z z) \times N \times \max_{i=1}^N (b_i - a_i + 1) \right)$ using dynamic programming. Please let me know if this will be useful and I would be happy to elaborate more.
– Sobi
Dec 19, 2015 at 22:34
• @becko, please let me know if you have questions about my response, or if it is not efficient enough.
– Sobi
Dec 23, 2015 at 3:10

We can compute the $p(z)$ using dynamic programming. Let's define $f(s, k)$ as follows:

$$f(s, k) = \sum_{x_k = a_k}^{b_k} \dots \sum_{x_N = a_N}^{b_N} \delta \left( s - \sum_{i=k}^N c_i x_i \right) p_k(x_k) \dots p_N(x_N). \tag{1}$$

Note that $p(z) \propto f(cz, 1)$, more precisely, $p(z) = \frac{f(cz, 1)}{\sum_{z'} f(cz', 1)}$.

# Computing $f(s, k)$ using dynamic programming

We can write the $f(s, k)$ recursively as follows: $$f(s, k) = \sum_{x_k = a_k}^{b_k} p_k(x_k) f(s - c_k x_k, k + 1). \tag{2}$$ Boundary cases include:

• If $s < 0$ then $f(s, k) = 0, \forall k$.
• If $s > 0$ then $f(s, k) = 0, \forall k > N$.
• $f(0, k) = \begin{cases} 1 & \text{ if } k = N+1 \\ 0 & \text{ otherwise } \end{cases}$.

Now think of $f(s, k)$ as a matrix $M$ of size $(c \max_z + 1) \times (N + 1)$ where $M_{s, k} = f(s, k)$. We can use Equation $(2)$, together with the boundary conditions, to compute all the entries in $M$. We can start from the entry at the top-right corner of the matrix (i.e. $M_{0, N}$) and move downwards to compute the first column. Then we can proceed to computing the values on the next column, starting from the top (i.e. $M_{0, N-1}$).

# Computational complexity

Entry $M_{s,k}$ of the matrix takes $O(b_k - a_k + 1)$ steps to compute (see Equation $(2)$) and the matrix has $(c \max_z + 1) \times (N + 1)$ entries. So, altogether, the order of complexity of the algorithm is $O(c \mathbf{z} N \mathbf{d})$ where $\mathbf{z} = \max_z z$ and $\mathbf{d} = \max_{i=1}^N \left( b_i - a_i + 1 \right)$.

# NOTE:

For simplicity of exposition, I am assuming all $a_i >= 0, 1 \leq i \leq N$. However, with a few modification, a similar idea can be applied when $a_i$s are allowed to be negative too.