CODEWITHSUNDEEP
rdata.table
Kristofersen
Kristofersen

Reputation: 2806

Carry data.table column values forward when conditions are met

I have a data.table with two columns. 

dt = data.table(a = c(0,0,-1,rep(0,3),-1,1), b = c(1,2,3,2,4,2,4,5))
> dt
    a b
1:  0 1
2:  0 2
3: -1 3
4:  0 2
5:  0 4
6:  0 2
7: -1 4
8:  1 5

What I need to happen is anytime column a == -1 I need the value in column b carried forward to the spot before the next row where column a == -1. If there are no more -1s then the value in column b needs to continue until the end of the data.table

This is the result I'm hoping for

    a b
1:  0 1
2:  0 2
3: -1 3
4:  0 3
5:  0 3
6:  0 3
7: -1 4
8:  1 4

Upvotes: 4

Views: 262

Answers (2)

989
989

Reputation: 12935

Maybe something like this in base R:

x <- c(which(dt==-1), nrow(dt)+1)
#[1] 3 7 9
dt[x[1]:nrow(dt),]$b <- rep(dt$b[head(x,-1)], diff(x))

#   a b
#1:  0 1
#2:  0 2
#3: -1 3
#4:  0 3
#5:  0 3
#6:  0 3
#7: -1 4
#8:  1 4

Upvotes: 2

Kristofersen
Kristofersen

Reputation: 2806

Alright, this wasn't as difficult as I originally thought. I can delete this question if necessary, but I haven't found anything similar on stackoverflow so I'll just post my solution for now.

There was an issue with the first solution. This one actually does what I expect it to, but I'm sure there is a much faster way to calculate this.

library(data.table)
dt = data.table(a = c(0,0,-1,rep(0,3),-1,1), b = c(1,2,3,2,4,2,4,5))

indices = which(dt$a == -1)
values = dt$b[indices]

dt[ , "tmp" := findInterval(1:nrow(dt), indices)]

dt$b = mapply(function(tmp, b){
                      if(tmp == 0){
                        return(b)
                      }else{
                        return(values[tmp])
                      }
                    }, dt$tmp, dt$b)

dt[ , "tmp" := NULL]

> dt
    a b
1:  0 1
2:  0 2
3: -1 3
4:  0 3
5:  0 3
6:  0 3
7: -1 4
8:  1 4

Better solution thanks to @Frank

dt[, tmp := cumsum(a==-1)][tmp > 0L, b := first(b), by=tmp][, tmp := NULL ]

Upvotes: 4

Related Questions