按组基于条件行每n行计算线性回归斜率的实现方法

现有如下名为df的DataFrame：

# A tibble: 14 × 5
   group date                indicator value diff_hours
   <chr> <dttm>              <lgl>     <dbl>      <dbl>
 1 A     2022-11-01 01:00:00 FALSE         2          4
 2 A     2022-11-01 02:00:00 FALSE         1          3
 3 A     2022-11-01 03:00:00 FALSE         4          2
 4 A     2022-11-01 04:00:00 FALSE         1          1
 5 A     2022-11-01 05:00:00 TRUE          3          0
 6 A     2022-11-01 06:00:00 FALSE         1          1
 7 A     2022-11-01 07:00:00 FALSE         3          2
 8 B     2022-11-01 01:00:00 FALSE         1          4
 9 B     2022-11-01 02:00:00 FALSE         2          3
10 B     2022-11-01 03:00:00 FALSE         3          2
11 B     2022-11-01 04:00:00 FALSE         1          1
12 B     2022-11-01 05:00:00 TRUE          4          0
13 B     2022-11-01 06:00:00 FALSE         1          1
14 B     2022-11-01 07:00:00 FALSE         5          2

需求 #

按组处理数据，以indicator == TRUE的行为分界点，将分界点前后的行分别按每n行一组，对每组计算线性回归lm(value ~ diff_hours)的斜率；同时将indicator == TRUE的行的斜率设为NA。当n=2时，期望输出如下df_desired：

# A tibble: 14 × 6
# Groups:   group [2]
   group date                indicator value diff_hours slope
   <chr> <dttm>              <lgl>     <dbl>      <dbl> <dbl>
 1 A     2022-11-01 01:00:00 FALSE         2          4     1
 2 A     2022-11-01 02:00:00 FALSE         1          3     1
 3 A     2022-11-01 03:00:00 FALSE         4          2     3
 4 A     2022-11-01 04:00:00 FALSE         1          1     3
 5 A     2022-11-01 05:00:00 TRUE          3          0    NA
 6 A     2022-11-01 06:00:00 FALSE         1          1     2
 7 A     2022-11-01 07:00:00 FALSE         3          2     2
 8 B     2022-11-01 01:00:00 FALSE         1          4    -1
 9 B     2022-11-01 02:00:00 FALSE         2          3    -1
10 B     2022-11-01 03:00:00 FALSE         3          2     2
11 B     2022-11-01 04:00:00 FALSE         1          1     2
12 B     2022-11-01 05:00:00 TRUE          4          0    NA
13 B     2022-11-01 06:00:00 FALSE         1          1     4
14 B     2022-11-01 07:00:00 FALSE         5          2     4

示例：第1、2行组成的组，通过lm(c(2,1)~c(4,3))计算得到斜率为1。

数据的dput代码 #

df <- structure(list(group = c("A", "A", "A", "A", "A", "A", "A", "B", 
"B", "B", "B", "B", "B", "B"), date = structure(c(1667260800, 
1667264400, 1667268000, 1667271600, 1667275200, 1667278800, 1667282400, 
1667260800, 1667264400, 1667268000, 1667271600, 1667275200, 1667278800, 
1667282400), class = c("POSIXct", "POSIXt"), tzone = ""), indicator = c(FALSE, 
FALSE, FALSE, FALSE, TRUE, FALSE, FALSE, FALSE, FALSE, FALSE, 
FALSE, TRUE, FALSE, FALSE), value = c(2, 1, 4, 1, 3, 1, 3, 1, 
2, 3, 1, 4, 1, 5), diff_hours = c(4, 3, 2, 1, 0, 1, 2, 4, 3, 
2, 1, 0, 1, 2)), class = c("grouped_df", "tbl_df", "tbl", "data.frame"
), row.names = c(NA, -14L), groups = structure(list(group = c("A", 
"B"), .rows = structure(list(1:7, 8:14), ptype = integer(0), class = c("vctrs_list_of", 
"vctrs_vctr", "list"))), class = c("tbl_df", "tbl", "data.frame"
), row.names = c(NA, -2L), .drop = TRUE))

df_desired <- structure(list(group = c("A", "A", "A", "A", "A", "A", "A", "B", 
"B", "B", "B", "B", "B", "B"), date = structure(c(1667260800, 
1667264400, 1667268000, 1667271600, 1667275200, 1667278800, 1667282400, 
1667260800, 1667264400, 1667268000, 1667271600, 1667275200, 1667278800, 
1667282400), class = c("POSIXct", "POSIXt"), tzone = ""), indicator = c(FALSE, 
FALSE, FALSE, FALSE, TRUE, FALSE, FALSE, FALSE, FALSE, FALSE, 
FALSE, TRUE, FALSE, FALSE), value = c(2, 1, 4, 1, 3, 1, 3, 1, 
2, 3, 1, 4, 1, 5), diff_hours = c(4, 3, 2, 1, 0, 1, 2, 4, 3, 
2, 1, 0, 1, 2), slope = c(1, 1, 3, 3, NA, 2, 2, -1, -1, 2, 2, 
NA, 4, 4)), row.names = c(NA, -14L), class = c("grouped_df", 
"tbl_df", "tbl", "data.frame"), groups = structure(list(group = c("A", 
"B"), .rows = structure(list(1:7, 8:14), ptype = integer(0), class = c("vctrs_list_of", 
"vctrs_vctr", "list"))), class = c("tbl_df", "tbl", "data.frame"
), row.names = c(NA, -2L), .drop = TRUE))

可以通过dplyr结合tidyr实现需求，步骤如下：

1. 按组创建分界分组：以indicator == TRUE为分界，将每组数据拆分成多个子段（分界点前后各为一个子段）
2. 在每个子段内，按每n行创建计算分组
3. 对每个计算分组拟合线性回归，提取斜率
4. 将斜率匹配回原始数据，同时将indicator == TRUE的行斜率设为NA

具体代码如下（以n=2为例）：

library(dplyr)
library(tidyr)

n <- 2

df_result <- df %>%
  group_by(group) %>%
  # 创建分界分组：indicator为TRUE时，分界分组号+1
  mutate(segment = cumsum(indicator)) %>%
  group_by(group, segment) %>%
  # 在每个分界段内，按每n行创建计算分组
  mutate(calc_group = (row_number() - 1) %/% n) %>%
  group_by(group, segment, calc_group) %>%
  # 计算斜率：如果分组内有至少2行（lm需要至少2个点），否则设为NA
  mutate(slope = if(n() >= 2) coef(lm(value ~ diff_hours))[["diff_hours"]] else NA) %>%
  # 将indicator为TRUE的行斜率设为NA
  mutate(slope = ifelse(indicator, NA, slope)) %>%
  ungroup() %>%
  # 移除中间分组列
  select(-segment, -calc_group)

# 验证结果是否与期望一致
all.equal(df_result, df_desired)

代码解释 #

• segment = cumsum(indicator)：通过累加indicator的逻辑值，将每组数据拆分为分界点前后的子段，比如indicator为TRUE的行所在位置会将数据分成前后两个子段
• calc_group = (row_number() - 1) %/% n：在每个子段内，将行号转换为每n行一组的分组编号，实现按n行分组
• coef(lm(value ~ diff_hours))[["diff_hours"]]：拟合线性回归后提取diff_hours对应的系数，即斜率
• ifelse(indicator, NA, slope)：将indicator为TRUE的行斜率强制设为NA

运行上述代码后，得到的df_result与df_desired完全一致。

内容的提问来源于stack exchange，提问作者Quinten