CODEWITHSUNDEEP
rdata.table
Henk
Henk

Reputation: 73

R add new column depending on values in a range in different columns

I have a data.table with two boolean columns A and B. I'd like to add a new boolean row C that depends on A and B, but I'm having trouble 'looking' in previous and upcomming rows.

I'd like to define C as follows. If there is a row with A=1 and atleast one B=1 in a range of three rows around then I'd like for C to become C=1 on the row where A=1 and C=0 on all the other rows in the range. Else C should be C=B.

In the case that two ranges overlap and both contain a B=1 then C should become C=1 on both rows where A=1 and C=0 on the others. For more clarification:

df <- data.table(A=c(0,0,0,1,0,0,0,0,0,0,0,1,1,0,0), 
                 B=c(0,1,0,0,0,1,0,1,1,0,0,0,0,0,1))

    A B                                        A B C
1:  0 0 #                                  1:  0 0 0
2:  0 1 #                                  2:  0 1 0
3:  0 0 #                                  3:  0 0 0
4:  1 0 # range of three                   4:  1 0 1
5:  0 0 #                                  5:  0 0 0
6:  0 1 #                                  6:  0 1 0
7:  0 0 #                                  7:  0 0 0
8:  0 1                                    8:  0 1 1 # C = B
9:  0 1 #                                  9:  0 1 0
10: 0 0 ##                                 10: 0 0 0
11: 0 0 ##                                 11: 0 0 0
12: 1 0 ## overlapping range of three      12: 1 0 1
13: 1 0 ##                                 13: 1 0 1
14: 0 0 ##                                 14: 0 0 0
15: 0 1 ##                                 15: 0 1 0

How do I go about doing this, I'm kind of clueless on this one.

Upvotes: 0

Views: 613

Answers (2)

IceCreamToucan
IceCreamToucan

Reputation: 28675

# Find ranges where A == 1
ind <- lapply(which(df$A == 1)
              , function(i){s <- i + -3:3; s[s %in% seq(nrow(df))]})
# Remove ranges with no B == 1
good <- sapply(ind, function(i) df[i, any(B == 1)])
ind  <- unique(unlist(ind[good]))
# Assign C as described
df[, C := B]
df[ind, C := as.numeric(A == 1)]
df
#     A B C
#  1: 0 0 0
#  2: 0 1 0
#  3: 0 0 0
#  4: 1 0 1
#  5: 0 0 0
#  6: 0 1 0
#  7: 0 0 0
#  8: 0 1 1
#  9: 0 1 0
# 10: 0 0 0
# 11: 0 0 0
# 12: 1 0 1
# 13: 1 0 1
# 14: 0 0 0
# 15: 0 1 0

Data used below. I changed your df definition to match the df displayed

df <- data.table(A=c(0,0,0,1,0,0,0,0,0,0,0,0,1,0,0), 
                 B=c(0,1,0,0,0,1,0,1,1,0,0,0,0,0,0))

df[12, A := 1]
df[15, B := 1]

df

#     A B
#  1: 0 0
#  2: 0 1
#  3: 0 0
#  4: 1 0
#  5: 0 0
#  6: 0 1
#  7: 0 0
#  8: 0 1
#  9: 0 1
# 10: 0 0
# 11: 0 0
# 12: 1 0
# 13: 1 0
# 14: 0 0
# 15: 0 1

Upvotes: 3

Melissa Key
Melissa Key

Reputation: 4551

Here's a solution based on the tidyverse suite of packages:

I defined 2 temporary variables - A1 determines whether A = 1 anywhere in the (row -3 : row + 3) window. C1 tests whether A = 1 and B = 1 anywhere on the window. 

library(tidyverse)
df %>% 
  mutate(
    A1 = (cumsum(lead(A, 3, default = 0)) - cumsum(dplyr::lag(A, 4, default = 0)) > 0),
    C1 = (A & dplyr::lead(cumsum(B), n = 3, default = 0) - dplyr::lag(cumsum(B), n = 4, default = 0)) * 1,
    C = ifelse(!A1, B, C1)
  ) %>%
  select(-A1, -C1)

Upvotes: 2

Related Questions