Assign progressive number if a condition is met - R

时间:2016-08-31 16:57:02

标签: r

How do I assign progressive numbers to a column every time a given condition is met in another one? Given this data:

a <- data.frame(var = c(1, 0, 0, 1, 4, 5, 6, 1, 7, 1, 1))

I would like to construct a column that is progressively augmented by 1 every time var == 1 and returns NAs for the rest. The new column should then be filled with:

1, NA, NA, 2, NA, NA, NA, 3, NA, 4, 5

I thought about ifelse but I didn't manage to make it work. Thanks!

3 个答案:

答案 0 :(得分:3)

You can use ifelse and cumsum:

a$newvar <- ifelse(a$var==1, cumsum(a$var==1), NA)
   var newvar
1    1      1
2    0     NA
3    0     NA
4    1      2
5    4     NA
6    5     NA
7    6     NA
8    1      3
9    7     NA
10   1      4
11   1      5

答案 1 :(得分:0)

We can also do this with a variation of cumsum

a$newVar <- with(a, cumsum(var ==1) * NA^(var!=1))
a$newVar
#[1]  1 NA NA  2 NA NA NA  3 NA  4  5

Or using data.table, we convert the 'data.frame' to 'data.table' (setDT(a)), specify the logical condition in 'i' (var == 1), and assign (:= it is efficient as it assigns in place) the cumsum of 'var' to 'newvar'. By default, the other elements in 'newvar' that do not correspond to the logical condition will be filled by NA.

library(data.table)
setDT(a)[var==1, newvar := cumsum(var)]
a
#    var newvar
# 1:   1      1
# 2:   0     NA
# 3:   0     NA
# 4:   1      2
# 5:   4     NA
# 6:   5     NA
# 7:   6     NA
# 8:   1      3
# 9:   7     NA
#10:   1      4
#11:   1      5

Or instead of cumsum we can use the sequence of rows

setDT(a)[var==1, newvar := seq_len(.N)]

答案 2 :(得分:0)

I'd just like to add one more thing for a general case in which if you want to do the same thing for 4 or 5 or any thing else

a <- data.frame(var = c(1, 0, 0, 1, 4, 5, 6, 1, 7, 1, 1))

a$New <- ifelse(a$var == 1,1,NA)
a$New[!is.na(a$New)] <- cumsum(a$New[!is.na(a$New)])

Output:

> print(a)
   var New
1    1   1
2    0  NA
3    0  NA
4    1   2
5    4  NA
6    5  NA
7    6  NA
8    1   3
9    7  NA
10   1   4
11   1   5