Как использовать cummean на основе условия в dplyr

Question

Я пытаюсь вычислить кумулятивное среднее значение столбца в моем фрейме данных, но я хочу, чтобы это делалось только на основе условия.

data <- data.frame(col1 = c("A","A","B","B","A"),
                   col2 = c(1,0,1,0,1))

Допустим, я хочу вычислить col3 как cummean из col2, но только если col1 = "A"

Ожидаемый результат должен быть

A 1 1.00
A 1 0.50
B 1 NA
B 0 NA
A 1 0.67

В идеале было бы неплохо, если бы значения NA приняли значение lag cummean по умолчанию. Что-то вроде

A 1 1.00
A 1 0.50
B 1 0.50
B 0 0.50
A 1 0.67

Я пытался mutate_if, но не мог понять правильный синтаксис. Возможно, ответ должен быть cummean внутри ifelse с обычным mutate. Не уверен, что это лучший способ подойти. Любая помощь очень ценится. Спасибо!

Answer 1

FWIW, я сам нашел решение

data %>%
  mutate(col3 = ifelse(col1 == A, cummean(col2), NA)) %>%
  tidyr::fill(col3, .direction = "up") ## for filling up the NA values with previous values 

Любой чище / лучше подходит?

Answer 2

Рассчитать совокупную сумму col2 условно, а затем разделить на совокупное число A в col1:

data %>% 
  mutate(cummean = { 
    isA <- col1 == 'A'
    cumsum(isA * col2) / cumsum(isA) 
  })

#  col1 col2   cummean
#1    A    1 1.0000000
#2    A    0 0.5000000
#3    B    1 0.5000000
#4    B    0 0.5000000
#5    A    1 0.6666667