如何用apply简洁实现数据框非NA值替换及final列生成？

初始数据 #

首先生成含NA值的原始数据框：

mydf=data.frame(ID=LETTERS[1:10], aaa=runif(10), bbb=runif(10), ccc=runif(10), ddd=runif(10))
mydf[c(1,4,5,7:10),2]=NA
mydf[c(1,2,4:8),3]=NA
mydf[c(3,4,6:10),4]=NA
mydf[c(1,3,4,6,9,10),5]=NA

需求说明 #

1. 非NA值替换：将aaa至ddd列中的非NA值替换为「列索引-1」，替换后的数据框格式如下：

> mydf
   ID       aaa       bbb        ccc       ddd
1   A        NA        NA          3        NA
2   B         1        NA          3         4
3   C         1         2         NA        NA
4   D        NA        NA         NA        NA
5   E        NA        NA          3         4
6   F         1        NA         NA        NA
7   G        NA        NA         NA         4
8   H        NA        NA         NA         4
9   I        NA         2         NA        NA
10  J        NA         2         NA        NA

2. 添加final列：按以下规则生成final列：
   • 该行全为NA时，值为0；
   • 该行非NA值数量大于1时，值为0；
   • 该行仅含1个非NA值时，值为该非NA值。
     最终结果格式如下：

> mydf
   ID       aaa       bbb        ccc       ddd       final
1   A        NA        NA          3        NA           3
2   B         1        NA          3         4           0
3   C         1         2         NA        NA           0
4   D        NA        NA         NA        NA           0
5   E        NA        NA          3         4           0
6   F         1        NA         NA        NA           1
7   G        NA        NA         NA         4           4
8   H        NA        NA         NA         4           4
9   I        NA         2         NA        NA           2
10  J        NA         2         NA        NA           2

要求使用apply系列函数实现，避免for循环。

解决方案代码 #

# 第一步：替换非NA值为对应列索引减1
mydf[, 2:5] = apply(mydf[, 2:5], 2, function(col) {
  # 获取当前列的索引并减1
  col_val = which(colnames(mydf) == colnames(col)) - 1
  # 将非NA值替换为col_val
  col[!is.na(col)] = col_val
  return(col)
})

# 第二步：生成final列
mydf$final = apply(mydf[, 2:5], 1, function(row) {
  # 提取该行非NA值
  non_na = row[!is.na(row)]
  len = length(non_na)
  # 根据非NA值数量返回对应结果
  if (len == 0 || len > 1) {
    return(0)
  } else {
    return(non_na)
  }
})

# 查看最终结果
print(mydf)

代码解释 #

1. 非NA值替换：

   • 使用apply(..., 2, ...)按列处理第2到第5列；
   • 对每一列，先获取其在数据框中的索引，计算出列索引-1的值，再替换该列所有非NA值。

2. final列生成：

   • 使用apply(..., 1, ...)按行处理第2到第5列；
   • 提取每行的非NA值，根据其数量判断返回0还是该非NA值，逻辑清晰且避免了循环。

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