Creating all multi indices of given length

Viewed 34

Given an integer K and a dimension d, how can I get all multi-indices alpha with length(alpha) = d and sum(alpha) = K in R?

Example: For K=3 and d=2 and if we organize the multi-indices in a list alphas we would get

alphas[[1]] = c(3,0)
alphas[[2]] = c(2,1)
alphas[[3]] = c(1,2)
alphas[[4]] = c(0,3)
1 Answers

Based on the description, we may use expand.grid on the sequence from '0' to 'k', after replicating the list 'd' times, then Filter only the rows having the sum as 'k'

f1 <- function(k, d) { 
 lapply(Filter(function(x) sum(x) == k, 
      asplit(expand.grid(rep(list(0:k), d)), 1)), unname)

}

-testing

> f1(3, 2)
[[1]]
[1] 3 0

[[2]]
[1] 2 1

[[3]]
[1] 1 2

[[4]]
[1] 0 3

Or slightly more faster would be to filter with rowSums

d1 <- expand.grid(rep(list(0:3), 2))
asplit(unname(d1)[rowSums(d1) == 3,], 1)

There is also a constraints based combinations functions in RcppAlgos

f2 <- function(k, d) {

 out <- RcppAlgos::comboGeneral(0:k, d, constraintFun = "sum", 
      comparisonFun = "==", limitConstraints = k)
 asplit(rbind(out, out[, ncol(out):1]), 1)

}

-testing

> f2(3, 2)
[[1]]
[1] 0 3

[[2]]
[1] 1 2

[[3]]
[1] 3 0

[[4]]
[1] 2 1

Or as @JosephWood mentioned permuteGeneral would be more adequate compared to comboGeneral

k <- 3
d <- 2
permuteGeneral(0:k, d, TRUE, constraintFun = "sum",
  comparisonFun = "==", limitConstraints = k)

Or with compositions

library(partitions)
asplit(compositions(k, d), 2)
Related