How do you sample groups in a data.table with a caveat

Question

This question is very similar to Sample random rows within each group in a data.table.

The difference is in a minor subtlety that I did not have enough reputation to discuss for that question itself.

Let's change Christopher Manning's initial data a little bit:

> DT = data.table(a=c(1,1,1,1:15,1,1), b=sample(1:1000,20))
> DT
     a   b
 1:  1 102
 2:  1   5
 3:  1 658
 4:  1 499
 5:  2 632
 6:  3 186
 7:  4 761
 8:  5 150
 9:  6 423
10:  7 832
11:  8 883
12:  9 247
13: 10 894
14: 11 141
15: 12 891
16: 13 488
17: 14 101
18: 15 677
19:  1 400
20:  1 467

If we tried the question's solution:

> DT[,.SD[sample(.N,3)],by = a]

Error in sample.int(x, size, replace, prob) : cannot take a sample larger than the population when 'replace = FALSE'

This is because there are values in column that only occur once. We cannot sample 3 times for values that occur less than three times without using replacement (which we do not want to do).

I am struggling to deal with this scenario. We want to sample 3 times when the number of occurrences is >= 3, but pull the number of occurrences if it is < 3. For example with our DT above we would want:

     a   b
 1:  1 102
 2:  1   5
 3:  1 658
 4:  2 632
 5:  3 186
 6:  4 761
 7:  5 150
 8:  6 423
 9:  7 832
10:  8 883
11:  9 247
12: 10 894
13: 11 141
14: 12 891
15: 13 488
16: 14 101
17: 15 677

Maybe a solution could involve sorting the data.table like this, then using rle() lengths to find out which n to use in the sample function above:

> DT <- DT[order(DT$a),]
> DT
     a   b
 1:  1 102
 2:  1   5
 3:  1 658
 4:  1 499
 5:  1 400
 6:  1 467
 7:  2 632
 8:  3 186
 9:  4 761
10:  5 150
11:  6 423
12:  7 832
13:  8 883
14:  9 247
15: 10 894
16: 11 141
17: 12 891
18: 13 488
19: 14 101
20: 15 677

> ifelse(rle(DT$a)$lengths >= 3, 3,rle(DT$a)$lengths)
> [1] 3 1 1 1 1 1 1 1 1 1 1 1 1 1 1

If we replace "3" with n, this will return how much we should sample from a=1, a=2, a=3... I have yet to find a way to incorporate this into a final solution. Any help would be appreciated!

Answer 1

I might be misunderstanding your question, but are you looking for something like this?

set.seed(123)
##
DT <- data.table(
  a=c(1,1,1,1:15,1,1), 
  b=sample(1:1000,20))
##
R> DT[,.SD[sample(.N,min(.N,3))],by = a]
     a   b
 1:  1 288
 2:  1 881
 3:  1 409
 4:  2 937
 5:  3  46
 6:  4 525
 7:  5 887
 8:  6 548
 9:  7 453
10:  8 948
11:  9 449
12: 10 670
13: 11 566
14: 12 102
15: 13 993
16: 14 243
17: 15  42

where we are drawing 3 samples from b for group a_i if a_i contains three or more values, else we draw only n values, where n (n < 3) is the size of group a_i.

Just for demonstration, here are the 6 possible values of b for a=1 that we are sampling from (assuming you use the same random seed as above):

R> DT[order(a)][1:6,]
   a   b
1: 1 288
2: 1 788
3: 1 409
4: 1 881
5: 1 323
6: 1 996