# Add a column to a dataframe based on a probability distribution

Suppose I have two dependent categorical attributes A, B. I have a dataframe X that lists probabilities (or expected counts) for all combinations of all categories of A and B. Let's assume two categories in each attribute, the dataframe could then look like this:

X <- data.frame(A = c(1,1,2,2), B = 1:2, PROB = c(0.1,0.3,0.2,0.4))
X


Output:

  A B PROB
1 1 1   .1
2 1 2   .3
3 2 1   .2
4 2 2   .4


Now I would like to add a B column to another dataframe Y that has only an A column. Y exists and has many more rows than X, and also additional columns. A synthetic Y could look like this:

Y <- data.frame(A=sample(2, size=10000, replace=TRUE),
B=NA,
C=sample(5, size=10000, replace=TRUE),
D=sample(3, size=10000, replace=TRUE))


The contents of the new B column should be sampled at random using the conditional probabilities for B given the value of the A column.

I am new to R, and I was wondering whether this operation can be handled more elegantly than writing a loop that iterates over all A categories. Perhaps a dedicated model class for which predict could be applied? Also, I lack the knowledge of statistical terminology -- I would appreciate any hint on how this kind of imputation (?) is referred to in the literature.

• PROB is the overall probability, it sums up to 1 in my toy dataframe. In Y$B, I want the value 2 to appear about three times as often as the value 1 in those rows that have A=1. For the rows with A=2, this ratio should be 2:1 (=0.4/0.2). – krlmlr Sep 22 '11 at 11:58 • OK, so in the final dataframe only these conditional amounts should count and PROB won't come out the same as here. That had me stumped for a bit because your toy Y has equal probability 1 and 2 for A. – John Sep 22 '11 at 14:22 ## 2 Answers It seems to me that if Y doesn't actually already exist you probably want to just make it from scratch like this. Y <- X[sample(1:nrow(X), size = 10000, replace = TRUE, prob = X$PROB),]
Y <- Y[,-3]


BTW, you can test to see that you got what you want with...

table(factor(Y$A):factor(Y$B))/nrow(Y)


If Y really does already exist then yes, you have to go through the unique values of A, but not all values of A.

b <- lapply(unique(X$A), function(x){ n <- nrow(Y[Y$A==x,])
df <- X[X$A == x,] Y$B[Y$A==x] <<- with( df, sample(B, n, replace = TRUE, prob = PROB) ) })  In that case b is garbage. A for loop version of this wouldn't even make the garbage b for(x in unique(X$A)){
n <- nrow(Y[Y$A == x,]) df <- X[X$A == x,]
Y$B[Y$A == x] <- with( df, sample(B, n, replace = TRUE, prob = PROB) )
}


And thus... this might be a case where using a loop is the cleanest way to go. :)

I did some quick time testing of these two methods and the plyr version. The fastest is the for loop, followed by lapply, followed by (far in the back), ddply (50% with given dataset). The cost gets bigger as Y grows in either size or complexity such that the for loop can get 2 to 5x the performance of ddply and it uses much less memory.

• +1 for the table test. Is there really no shortcut for this assignment in R? I was thinking about a dedicated model class for which predict could be applied. – krlmlr Sep 22 '11 at 12:07
• predict couldn't work since it doesn't give random outcomes but fixed ones. – John Sep 22 '11 at 14:23

Use ddply (from plyr package) first to create the conditional probabilities:

Xcond<-ddply(X, .(A),
.fun=function(curdfr){
curdfr$CPROB<-curdfr$PROB/sum(curdfr$PROB);curdfr})  Now, given Y with column A, you can get the result you want e.g. by (if the order is not of importance to you): ddply(Y, .(A), .fun=function(curdfr){ curA<-curdfr[1,"A"] curOrg<-Xcond[Xcond$A==curA,]
curN<-nrow(curdfr)
curB<-sample(curOrg$B, size=nrow(curdfr), replace=TRUE, prob=curOrg$CPROB)
curdfr$B<-curB curdfr })  No doubt there are more elegant/efficient solutions. Read its help if you don't understand the ddply call - if that fails, get back to us. Edit as per @John's suggestion: X can be used immediately instead of Xcond: ddply(Y, .(A), .fun=function(curdfr){ curA<-curdfr[1,"A"] curOrg<-X[X$A==curA,]
curN<-nrow(curdfr)
curB<-sample(curOrg$B, size=nrow(curdfr), replace=TRUE, prob=curOrg$PROB)
curdfr\$B<-curB
curdfr
})

• +1 Thank you, that's definitely more elegant (and should be faster, too!) than looping through categories. – krlmlr Sep 22 '11 at 12:43
• Actually, that's less efficient than looping through categories. – John Sep 22 '11 at 13:23
• The first run to modify X is unnecessary. It will work fine with the original X. The values in prob for sample don't have to add to 1.0. The relative amounts will work. – John Sep 22 '11 at 13:25