Question
Iterate over indices in multinomial expansion
Context: In the multinomial expansion, the valid indices [k1, k2, ... kr] for each term fulfill the equation n = k1 + k2 + ... + kr. Given r and n, I'm trying to iterate over all terms. For that, I need a way to get every valid index set.
Example for r = 2 and n = 3:
[0 3], [1 2], [2 1], [3 0]
I tried to treat the valid arrays as numbers in base n+1, where I increment the number in the base by 1 and then check if the sum of all digits equals n.
To save time and memory, my function doesn't take r and n as inputs but just a previous valid index set.
function out_set = getNextExponent(inp_set)
% inp_set is a valid exponent
% returns the next valid exponent set of a multinomial expansion
% Examples:
% getNextExponent([0, 0, 3]) = [0, 1, 2]
% getNextExponent([0, 2, 1]) = [1, 1, 1]
% getNextExponent([3, 0, 0]) = error
% disp(getNextExponent([1, 0, 3, 5, 0])) = [1, 0, 4, 0, 4]
% disp(getNextExponent([100, 29, 29, 0, 0, 0])) = [100, 30, 0, 0, 0, 28]
% Get the parameters
n = sum(inp_set);
base = n+1; % exponent list is a number in base n+1
% Check if it is the last valid index set
out_set = zeros(size(inp_set));
out_set(1) = n;
if isequal(inp_set, out_set)
error('Last exponent already reached')
end
% Loop until next valid set is found
while true
% Add one to the set
carry = 1;
index = length(inp_set);
while carry == 1
if inp_set(index)+1 < base
inp_set(index) = inp_set(index)+1;
carry = 0;
else
inp_set(index) = 0;
index = index-1;
end
end
% Is the next valid exponent reached?
if sum(inp_set) == n
break;
end
end
out_set = inp_set;
end
However, this is extremely inefficient. Are there better ways to iterate over each valid set?
Edit: I don't know if it's helpful, but this is the same problem as putting n identical objects in r boxes. So the number of possible arrays is always (n+r-1 choose r-1).