How can I select half of the counts in R?

Viewed 85

I have a dataset with 2 columns like in the screenshot below: Personal name, Company name

> dput(df)
structure(list(Name = c("ABC", "BCD", "CDE", "DEF", "EFG", "FGH", 
"GHI", "HIJ", "IJK", "JKL"), Company = c("A", "A", "A", "A", 
"B", "B", "C", "C", "C", "C")), class = "data.frame", row.names = c(NA, 
-10L))

> df
   Name Company
1   ABC       A
2   BCD       A
3   CDE       A
4   DEF       A
5   EFG       B
6   FGH       B
7   GHI       C
8   HIJ       C
9   IJK       C
10  JKL       C

How can I split the dataset into 2 subsets, each of the subsets has the same amount of people from the same company?

For example, in total, there are 4 people in Company A. Within Company A group, there are two subsets, i.e.,2 people in subset1, the other 2 in subset 2. similarly for companies B, C.

For an odd number of people in the same Campany, the tie can be break by randomly selecting a subset.

Thanks!

3 Answers

try this tidyverse approach to get a RANDOM division of data into two sub-sets.

> df
   Name Company
1   ABC       A
2   BCD       A
3   CDE       A
4   DEF       A
5   EFG       B
6   FGH       B
7   GHI       C
8   HIJ       C
9   IJK       C
10  JKL       C

APPROACH-1 If you want results in a list with two separate dataframes

library(tidyverse)
set.seed(123) # set the seed to replicate the result lateron

df %>% mutate(id = row_number()) %>% group_by(Company) %>% mutate(n = n()) %>%
  sample_n(n*0.5) %>% mutate(d = 1) %>% ungroup() %>% select(id, d) %>% 
  right_join(df %>% mutate(id = row_number()), by = "id") %>%
  mutate(d = if_else(is.na(d), 2, d)) %>% group_split(d)


<list_of<
  tbl_df<
    id     : integer
    d      : double
    Name   : character
    Company: character
  >
>[2]>
[[1]]
# A tibble: 5 x 4
     id     d Name  Company
  <int> <dbl> <chr> <chr>  
1     3     1 CDE   A      
2     4     1 DEF   A      
3     5     1 EFG   B      
4     8     1 HIJ   C      
5     9     1 IJK   C      

[[2]]
# A tibble: 5 x 4
     id     d Name  Company
  <int> <dbl> <chr> <chr>  
1     1     2 ABC   A      
2     2     2 BCD   A      
3     6     2 FGH   B      
4     7     2 GHI   C      
5    10     2 JKL   C 

When you want results in two separate csv files.

set.seed(123)
df %>% mutate(id = row_number()) %>% group_by(Company) %>% mutate(n = n()) %>%
  sample_n(n*0.5) %>% mutate(d = 1) %>% ungroup() %>% select(id, d) %>% 
  right_join(df %>% mutate(id = row_number()), by = "id") %>%
  mutate(d = if_else(is.na(d), 2, d)) %>% group_split(d) %>% map2(c("file1.csv", "file2.csv"), ~ write.csv(.x, .y))

Above code will write 2 csv files in your getwd() directory.

APPROACH -2 Also you can do

set.seed(123)
df1 <- df %>% mutate(id = row_number()) %>% group_by(Company) %>% mutate(n = n()) %>%
  sample_n(n*0.5) %>% mutate(set_no = "Set1") %>% ungroup() %>% select(-n, -id)


df2 <- df %>% mutate(id = row_number()) %>% filter(!id %in% df1$id) %>% 
  mutate(set_no = "Set2") %>% select(-id)

Writing results in csv

write.csv(df1, "subset1.csv")
write.csv(df2, "subset2.csv")

Column id and d are helper columns. id columns will ensure that no row/record has been left in case there are odd entries in any group. d will help to identify the sub-group number.

In case any company has odd number of records, subsets may have different number of rows therein and not equal.

You could calculate modulo 2 of sequences of the subsets using ave.

r <- transform(d, sub=ave(company, company, FUN=function(x) 
  paste0("sub", seq_along(x) %% 2 + 1)))
r
#       name company  sub
# 1     alex       A sub2
# 2    malan       A sub1
# 3  matteis       A sub2
# 4  fenwick       A sub1
# 5  nicolas       B sub2
# 6   cleary       B sub1
# 7      fin       C sub2
# 8    stijn       C sub1
# 9  antoine       C sub2
# 10     fin       C sub1
# 11   stijn       C sub2
# 12 antoine       C sub1

Check:

with(r, table(company, sub))
# company sub1 sub2
#       A    2    2
#       B    1    1
#       C    1    2

Actually, if company is ordered as they're in the example, you could simply do:

d$sub <- NA
d$sub[] <- paste0("sub", 1:2)  ## throws a warning when `nrow` is uneven
d
#      name company  sub
# 1    alex       A sub1
# 2   malan       A sub2
# 3 matteis       A sub1
# 4 fenwick       A sub2
# 5 nicolas       B sub1
# 6  cleary       B sub2
# 7     fin       C sub1
# 8   stijn       C sub2
# 9 antoine       C sub1

Data:

d <- structure(list(name = c("alex", "malan", "matteis", "fenwick", 
"nicolas", "cleary", "fin", "stijn", "antoine"), company = c("A", 
"A", "A", "A", "B", "B", "C", "C", "C")), class = "data.frame", row.names = c(NA, 
-9L))

You can use ave and sample like:

set.seed(42)
df$su <- ave(df$Company, df$Company, FUN=function(x)
 sample(rep(1:2, ceiling(length(x)/2)), length(x)))
df
#   Name Company su
#1   ABC       A  1
#2   BCD       A  2
#3   CDE       A  1
#4   DEF       A  2
#5   EFG       B  2
#6   FGH       B  1
#7   GHI       C  2
#8   HIJ       C  2
#9   IJK       C  1
#10  JKL       C  1
Related