Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Avoiding writing a long if-else statement in R

Tags:

r

if-statement

I have run into a situation where I have a data like this:

df <- data.frame(id = 1:1000, 
                   x = sample(0:30, 1000, replace = T), 
                   y = sample(50:10000, 1000, replace = T))

I want to assign another column called z based on multiple conditions i.e.

if x <= 5 & y <= 100, z = 1
if x > 5 & x <= 10 & y <= 100, z = 2
if x > 10 & x <= 12 & y <= 100, z = 3
if x > 12 & x <= 20 &  y <= 100, z = 4
if x > 20 & x <= 30 &  y <= 100, z = 5
if x <= 5 & y > 100 & y <= 1000, z = 6
if x > 5 & x <= 10 & y > 100 & y <= 1000 z = 7
if x > 10 & x <= 12 & y > 100 & y <= 1000, z = 8
if x > 12 & x <= 20 & y > 100 & y <= 1000, z = 9
if x > 20 & x <= 30 & y > 100 & y <= 1000, z = 10
.
.
.

and so. I hope you get the drift.

The obvious solution for me to do is this to write a long ifelse statement something like this;

df %>% mutate(z = ifelse(x <= 5 & y <= 100, 1, 
                  ifelse(x > 5 & x <= 10 & y <= 100, 2,
                  ifelse(x > 10 & x <= 12 & y <= 100, 3))),
          ........... and son on)

You would find that such scripts can be endlessly long and I wondered if there are other ways to achieve this without writing the long ifelse statement.

like image 610
89_Simple Avatar asked Mar 04 '23 16:03

89_Simple


2 Answers

If there is a pattern in the if else statements, we can create the set of expressions beforehand and use !!! to unqoute and splice them into arguments to case_when:

x_gt_cond <- rep(c(-Inf, 5, 10, 12, 20), 2)
x_le_cond <- rep(c(5, 10, 12, 20 ,30), 2)
y_gt_cond <- rep(c(-Inf, 100), each = 5)
y_le_cond <- rep(c(100, 1000), each = 5)
z <- 1:10
cases <- paste("x > ", x_gt_cond, "& x <= ", x_le_cond, 
               "& y > ", y_gt_cond, "& y <= ", y_le_cond, "~ ", z)

library(dplyr)
library(rlang)
df %>%
  mutate(z = case_when(!!!parse_exprs(cases)))

The trick is to use -Inf and Inf for the lower and upper bounds so that you have balanced conditions for x and y. What's elegant about this solution is that you can add more conditions simply by altering the _cond vectors.

Output:

> cases
 [1] "x >  -Inf & x <=  5 & y >  -Inf & y <=  100 ~  1"
 [2] "x >  5 & x <=  10 & y >  -Inf & y <=  100 ~  2"  
 [3] "x >  10 & x <=  12 & y >  -Inf & y <=  100 ~  3" 
 [4] "x >  12 & x <=  20 & y >  -Inf & y <=  100 ~  4" 
 [5] "x >  20 & x <=  30 & y >  -Inf & y <=  100 ~  5" 
 [6] "x >  -Inf & x <=  5 & y >  100 & y <=  1000 ~  6"
 [7] "x >  5 & x <=  10 & y >  100 & y <=  1000 ~  7"  
 [8] "x >  10 & x <=  12 & y >  100 & y <=  1000 ~  8" 
 [9] "x >  12 & x <=  20 & y >  100 & y <=  1000 ~  9" 
[10] "x >  20 & x <=  30 & y >  100 & y <=  1000 ~  10"

       id  x    y  z
1       1 13 8440 NA
2       2  3 1467 NA
3       3  5 2699 NA
4       4 24 5286 NA
5       5  5 2378 NA
6       6 16  268  9
7       7 19 2910 NA
8       8 19  706  9
9       9 24 6212 NA
10     10  7 6026 NA
...
like image 133
acylam Avatar answered Mar 15 '23 10:03

acylam


It sounds like the case_when function in dplyr is what you're looking for. In your case, it might look something like this.

df %>% mutate(z = case_when(
   x <= 5 & y <= 100 ~ 1,
   x > 5 & x <= 10 & y <= 100 ~ 2,
   x > 10 & x <=12 & y <= 100 ~ 3
  )
 )

edit: Changed answer to reflect that case_when is in the dplyr package. Thanks for comments below.

like image 29
Ekholme Avatar answered Mar 15 '23 10:03

Ekholme