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)
CodePudding user response:
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)
