Saurabh
Saurabh

Reputation: 1626

R How to replace elements of a group only when all elements are zero

I have the data.table dt to which I want to add a new column new_col only when all elements in the group are zero.

library(data.table)
a <- c(1, 2, 3, 0, 0, 0, 0, 1)
group <- c("a", "a", "a", "b", "b", "c", "c", "c")
dt = data.table(a, group)
dt[]
   a group
1: 1     a
2: 2     a
3: 3     a
4: 0     b
5: 0     b
6: 0     c
7: 0     c
8: 1     c

Following is the code I tied.

dt[, new_col := ifelse(all(a == 0), NA, a), by = group]
dt[]

The output I am getting is

   a group new_col
1: 1     a       1
2: 2     a       1
3: 3     a       1
4: 0     b      NA
5: 0     b      NA
6: 0     c       0
7: 0     c       0
8: 1     c       0

The expected output is

   a group new_col
1: 1     a       1
2: 2     a       2
3: 3     a       3
4: 0     b      NA
5: 0     b      NA
6: 0     c       0
7: 0     c       0
8: 1     c       1

Can someone please point out what I am doing wrong?

Upvotes: 1

Views: 104

Answers (1)

akrun
akrun

Reputation: 887691

Here, we can use if/else as ifelse requires all arguments to be same length all(a == 0) is of length 1, along with the 'yes' but 'no' is not of length 1, which leads to recycling

dt[, new_col := if(all(a == 0)) NA else a, by = group]

-output

dt
   a group new_col
1: 1     a       1
2: 2     a       2
3: 3     a       3
4: 0     b      NA
5: 0     b      NA
6: 0     c       0
7: 0     c       0
8: 1     c       1

Or may also use .I to get the row index of the negated expression (or use any), extract ($V1) and assign those elements from 'a' to new_col. By default, the other elements will be NA

dt[dt[, .I[!all(a == 0)], group]$V1, new_col := a]

Upvotes: 2

Related Questions