用 R data.table 中的序列填充 NA 值

Question

我有一个数据 table 如下所示。请注意，当 vals 为 0 且其他地方缺失时，标志为 1。

dt <- data.table(vals = c(0,2,4,1,0,4,3,0,3,4))
dt[vals == 0, flag := 1]

 > dt
    vals flag
 1:    0    1
 2:    2   NA
 3:    4   NA
 4:    1   NA
 5:    0    1
 6:    4   NA
 7:    3   NA
 8:    0    1
 9:    3   NA
10:    4   NA

我希望输出看起来像下面的 seq 列。也就是说，当 vals 为 0 时，该列需要包含一组从 1 开始的序列，并在 vals 为 0 时向上计数直到下一行。 flag 只有在有助于实现所描述的目标时才有用。

 > dt
    vals  seq
 1:    0    1
 2:    2    2
 3:    4    3
 4:    1    4
 5:    0    1
 6:    4    2
 7:    3    3
 8:    0    1
 9:    3    3
10:    4    3

本来想用cumsum()的，但是想不出怎么用才有效。

我目前的解决方案非常丑陋。

dt <- data.table(vals = c(0,2,4,1,0,4,3,0,3,4))
dt[vals == 0, flag := 1]
dt[, flag_rleid := rleid(flag)]

# group on the flag_rleid column
dt[, flag_seq := seq_len(.N), by = flag_rleid]
# hideous subsetting to avoid incrementing the first appearance of a 1
dt[vals != 0, flag_seq := flag_seq + 1]

# flag_seq is the desired column
> dt
    vals flag flag_rleid flag_seq
 1:    0    1          1        1
 2:    2   NA          2        2
 3:    4   NA          2        3
 4:    1   NA          2        4
 5:    0    1          3        1
 6:    4   NA          4        2
 7:    3   NA          4        3
 8:    0    1          5        1
 9:    3   NA          6        2
10:    4   NA          6        3

感谢任何改进。

Answer 1

我们可以使用带有 cumsum 的逻辑索引来创建分组变量，然后基于它我们得到序列列

dt[, flag_seq := seq_len(.N), cumsum(vals ==0)]
dt
#    vals flag flag_seq
# 1:    0    1        1
# 2:    2   NA        2
# 3:    4   NA        3
# 4:    1   NA        4
# 5:    0    1        1
# 6:    4   NA        2
# 7:    3   NA        3
# 8:    0    1        1
# 9:    3   NA        2
#10:    4   NA        3

用 R data.table 中的序列填充 NA 值

Filling NA values with a sequence in R data.table

r

sequence

data.table