基于0/1/NA矩阵生成干预时间矩阵的代码修复请求

现有如下R语言的data.frame：

structure(list(`1` = c(0, 0, NA, NA, NA, NA, 0, 0, NA, NA, NA, 
1, NA, NA, NA), `2` = c(1, 0, NA, NA, NA, NA, NA, 0, NA, NA, 
NA, 1, NA, NA, NA), `4` = c(NA, NA, 0, 1, 1, 0, NA, NA, 0, 1, 
1, NA, 1, 0, 0), `5` = c(NA, NA, 0, 1, 1, 0, NA, NA, 1, 1, NA, 
NA, 1, 0, 1), `6` = c(NA, NA, 0, 1, 1, 0, NA, NA, 1, 0, NA, NA, 
1, 0, NA), `7` = c(NA, NA, NA, 1, 1, 0, NA, NA, 0, 1, NA, NA, 
1, 0, NA), `8` = c(NA, NA, NA, 1, 0, 0, NA, NA, 1, 0, NA, NA, 
1, 0, NA)), row.names = c(NA, 15L), class = "data.frame")

需要生成一个符合以下规则的新矩阵：

• 每行中若出现1，则该1之后的所有非NA元素都设为1；
• NA值保持不变。

示例：输入行0,NA,NA,1,0,NA,1，输出应为0,NA,NA,1,1,NA,1。

以下是编写的错误代码：

for(i in 1:nrow(mat)){
  for(j in 1:7){
    if(mat[i,j]==0){
      next }else{
    if(mat[i,j]==1){
        mat[i,j:7]<-1
      }else{
        if(is.na(mat[i,j])){
    mat[i,j]<-NA
      }}}
  }
  
}

原代码问题分析 #

1. 错误覆盖NA值：原代码中遇到1时，直接将j:7区间所有元素赋值为1，会把该区间内的NA也改成1，违反了NA保持不变的规则。
2. 循环逻辑冗余：未标记是否已遇到过1，导致后续循环重复处理已修改的元素，且没有针对NA做正确的跳过处理。

修复后的循环代码 #

# 先将数据赋值给mat（假设原始数据已存入该变量）
mat <- structure(list(`1` = c(0, 0, NA, NA, NA, NA, 0, 0, NA, NA, NA, 
1, NA, NA, NA), `2` = c(1, 0, NA, NA, NA, NA, NA, 0, NA, NA, 
NA, 1, NA, NA, NA), `4` = c(NA, NA, 0, 1, 1, 0, NA, NA, 0, 1, 
1, NA, 1, 0, 0), `5` = c(NA, NA, 0, 1, 1, 0, NA, NA, 1, 1, NA, 
NA, 1, 0, 1), `6` = c(NA, NA, 0, 1, 1, 0, NA, NA, 1, 0, NA, NA, 
1, 0, NA), `7` = c(NA, NA, NA, 1, 1, 0, NA, NA, 0, 1, NA, NA, 
1, 0, NA), `8` = c(NA, NA, NA, 1, 0, 0, NA, NA, 1, 0, NA, NA, 
1, 0, NA)), row.names = c(NA, 15L), class = "data.frame")

for(i in 1:nrow(mat)){
  encountered_one <- FALSE  # 标记是否已遇到第一个1
  for(j in 1:ncol(mat)){
    if(is.na(mat[i, j])){
      next  # NA值直接跳过，保持原样
    }
    if(encountered_one){
      mat[i, j] <- 1  # 已遇到1，非NA元素设为1
    } else {
      if(mat[i, j] == 1){
        encountered_one <- TRUE  # 标记已遇到第一个1，当前值保持1
      }
      # 若为0则保持原样，继续遍历
    }
  }
}

向量化实现方式（更高效） #

如果数据量较大，推荐使用apply函数实现向量化处理，避免循环：

process_row <- function(row){
  first_one_pos <- min(which(row == 1), Inf)
  if(is.finite(first_one_pos)){
    # 从第一个1的位置开始，将非NA元素设为1
    row[first_one_pos:length(row)][!is.na(row[first_one_pos:length(row)])] <- 1
  }
  return(row)
}

mat_processed <- t(apply(mat, 1, process_row))
# 转换回data.frame格式
mat_processed <- as.data.frame(mat_processed)

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