如何在dplyr中基于index列值提取分组后反转排序的非零value值？

数据准备 #

df <- data.frame(group = rep(c("A", "B"), each = 10),
                 value = c(0, 0, 11, 5, 9, 8, 0, 6, 0, 9, 
                           4, 0, 0, 0, 18, 1, 1, 0, 3, 6), 
                 index = c(0, 4, 3, 3, 2, 6, 0, 1, 0, 5, 
                           3, 0, 6, 0, 2, 4, 5, 0, 2, 1))

需求说明 #

需要按group分组添加新列column_wanted，规则如下：

• 当index列值为0时，新列值为0；
• 当index列值非零时，以该值为索引n，提取对应组中value列非零值经反转排序后的第n个值。

期望输出：

group value index column_wanted
1      A      0     0             0
2      A      0     4             8
3      A     11     3             9
4      A      5     3             9
5      A      9     2             9
6      A      8     6             5
7      A      0     0             0
8      A      6     1            11
9      A      0     0             0
10     A      9     5             6
11     B      4     3             4
12     B      0     0             0
13     B      0     6             1
14     B      0     0             0
15     B     18     2             6
16     B      1     4             3
17     B      1     5             1
18     B      0     0             0
19     B      3     2             6
20     B      6     1            18

当前实现方法（拆分合并） #

通过拆分数据为index=0和index≠0两个子集，分别处理后合并，能得到正确结果，但面对大量分组和观测时不够高效：

# 处理非零index的子集
df_no0 <- df %>% 
  filter(index != 0) %>%
  group_by(group) %>% 
  mutate(correct_col = rev(sort(value))[index])

# 处理零index的子集
df_just0 <- df %>% 
  filter(index == 0) %>%
  mutate(correct_col = 0)

# 合并结果
df_final <- bind_rows(df_no0, df_just0) %>% arrange(group)

尝试的错误方法及问题 #

尝试直接在mutate中用ifelse实现，但结果不符合预期：

df %>%
  group_by(group) %>%
  mutate(wrong_column = ifelse(index == 0, 0, 
            rev(sort(value[value != 0]))[index[index != 0]]))

错误原因：index[index != 0]提取的是当前组内非零的index值，得到的是一个长度短于原组的向量，ifelse会循环填充该向量，导致索引错位，结果错误。

另外尝试case_when时因左右侧长度不匹配报错，也是同样的逻辑问题。

简洁解决方案 #

在分组后先预生成每组的非零值反转排序向量，再用索引提取，避免拆分合并：

方法1：使用ifelse #

library(dplyr)

df_result <- df %>%
  group_by(group) %>%
  mutate(
    # 预存每组非零value反转排序后的向量，用list保存避免长度不匹配
    sorted_rev_nonzero = list(rev(sort(value[value != 0]))),
    column_wanted = ifelse(
      index == 0,
      0,
      # 从预存的向量中按index提取对应元素
      sorted_rev_nonzero[[1]][index]
    )
  ) %>%
  select(-sorted_rev_nonzero)  # 移除辅助列

# 查看结果
df_result

方法2：使用case_when #

df_result <- df %>%
  group_by(group) %>%
  mutate(
    sorted_rev_nonzero = list(rev(sort(value[value != 0]))),
    column_wanted = case_when(
      index == 0 ~ 0,
      TRUE ~ sorted_rev_nonzero[[1]][index]
    )
  ) %>%
  select(-sorted_rev_nonzero)

该方案无需拆分数据，直接在分组内完成计算，适合处理大规模数据，结果与期望输出一致。

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