状態を持つループ処理を accumulate() でシンプルに書く

R言語のコミュニティ https://r-wakalang.slack.com で回答したのでメモ。

質問はこんな感じ（意訳しています）。

次のようなデータを以下のルールで処理したい。データを上から下に見ていき、 (1) before に TRUE が出たら、それ以降は after を TRUE にする。 (2) ただし、condition が FALSE になったら after を FALSE にして状態をリセットする。これを、for を使わないやり方で書きたい。（データにすでにある after は答えあわせ用）

   before condition after
 1 FALSE  FALSE     FALSE
 2 FALSE  TRUE      FALSE
 3 TRUE   TRUE      TRUE 
 4 FALSE  TRUE      TRUE 
 5 FALSE  TRUE      TRUE 
 6 FALSE  FALSE     FALSE
 7 FALSE  TRUE      FALSE
 8 FALSE  TRUE      FALSE
 9 TRUE   TRUE      TRUE 
10 FALSE  TRUE      TRUE

以下、データは次のコードで生成したものを使う（元の質問より簡略化しています）。

library(tidyverse)

# データの準備
d <- tibble(
  before    = c(F,F,T,F,F,F,F,F,T,F),
  condition = c(F,T,T,T,T,F,T,T,T,T), 
  after     = c(F,F,T,T,T,F,F,F,T,T),
)

for を使った書き方

これは、状態 (state) を持つ繰り返し処理であり、for を使って書ける。

データを1行ずつ上から処理する。初期値は before = FALSE, condition = FALSE でこのときの出力 after は FALSE
初期状態 (state = 0) で before が TRUE になると after を TRUE にして次の状態 (state = 1) に移る。before が FALSE ならば after は FALSE のままで状態も変わらない
次の状態 (state = 1) で condition が FALSE になると after を FALSE にして初期状態に戻る。condition が TRUE ならば after は TRUE のままで状態も変わらない

state <- 0  # 状態を初期化
for (i in seq_len(nrow(d))) {
  row <- d[i, ]  # 一行ずつ取得
  if (state == 0) {
    if (row$before == TRUE) {
      d[i, "after2"] <- TRUE
      state <- 1  # 次の状態に移る
    } else {
      d[i, "after2"] <- FALSE
    }
  } else {  # state == 1
    if (row$condition == FALSE) {
      d[i, "after2"] <- FALSE
      state <- 0  # 初期状態に戻る
    } else {
      d[i, "after2"] <- TRUE
    }
  }
}

結果は d の after2 に格納される。期待する結果である after と一致する。

   before condition after after2
 1 FALSE  FALSE     FALSE FALSE 
 2 FALSE  TRUE      FALSE FALSE 
 3 TRUE   TRUE      TRUE  TRUE  
 4 FALSE  TRUE      TRUE  TRUE  
 5 FALSE  TRUE      TRUE  TRUE  
 6 FALSE  FALSE     FALSE FALSE 
 7 FALSE  TRUE      FALSE FALSE 
 8 FALSE  TRUE      FALSE FALSE 
 9 TRUE   TRUE      TRUE  TRUE  
10 FALSE  TRUE      TRUE  TRUE

accumulate() の使い方

これを for などのループを使わずに書くにはどうすればよいだろうか？

R言語の通常のベクトル化された関数では、状態を持つ処理を行うのは困難である。

そこで、accumulate() 関数を使う。この関数は、ベクトルを1つずつ処理していく関数であり、その際に前の処理の結果を使うことができる（基本関数の Reduce() と同様）。

状態の扱いが問題だったので、状態をいったん出力し、その状態を使って after を求めるという方針を取ろう。まずは、前の状態を入力にとり、次の状態を結果として返す関数を作成する（上記の for の中身から state に関する部分を抜き出せば簡単にできる）。

compute_state <- function(state, before, condition) {
  if (state == 0) {
    if (before == TRUE) {
      state <- 1
    }
  } else {  # state == 1
    if (condition == FALSE) {
      state <- 0
    }
  }
  state
}

accumulate2() 関数を使って状態を計算するには次のように書く（accumulate() は1変数を受け取り、accumulate2() は2変数を受け取る）。

state_vec <- accumulate2(d$before, d$condition, compute_state, .init = 0)
# 0 0 0 1 1 1 0 0 0 1 1

length(state_vec)
# 11

ただし、得られる結果はデータの行数より1つ多い。これは初期値 (.init = 0) が先頭に含まれるためである。

そのため、現在の状態を得るには最後の値を取り除けばよい。

current_state_vec <- state_vec |> head(-1)
# 0 0 0 1 1 1 0 0 0 1

ちなみに、次の状態を得たい場合は先頭を取り除く。

next_state_vec <- state_vec |> tail(-1)
# 0 0 1 1 1 0 0 0 1 1

accumulate() を使った書き方

現在の状態 state を得ることができた。出力 after は state, before, condition を使って通常のベクトル化された関数で計算できる。

state が 0 のとき、before が TRUE の場合のみ after = TRUE
state が 1 のとき、condition が FALSE の場合のみ after = FALSE、すなわち condition が TRUE の場合のみ after = TRUE

d |>
  mutate(state = accumulate2(before, condition, compute_state, .init = 0) |> head(-1)) |>
  mutate(after3 = (state == 0 & before) | (state == 1 & condition))

   before condition after after2 state after3
 1 FALSE  FALSE     FALSE FALSE      0 FALSE 
 2 FALSE  TRUE      FALSE FALSE      0 FALSE 
 3 TRUE   TRUE      TRUE  TRUE       0 TRUE  
 4 FALSE  TRUE      TRUE  TRUE       1 TRUE  
 5 FALSE  TRUE      TRUE  TRUE       1 TRUE  
 6 FALSE  FALSE     FALSE FALSE      1 FALSE 
 7 FALSE  TRUE      FALSE FALSE      0 FALSE 
 8 FALSE  TRUE      FALSE FALSE      0 FALSE 
 9 TRUE   TRUE      TRUE  TRUE       0 TRUE  
10 FALSE  TRUE      TRUE  TRUE       1 TRUE

非常にシンプルに書くことができた。

余談

ここからは余談であるが、この処理はもっとシンプルにできる。

まず、状態を計算する関数 compute_state() は次のように書ける。

compute_state <- function(state, before, condition) {
  if (state == 0) {
    before
  } else {  # state == 1
    condition
  }
}

さらにシンプルにすると次のようになる。

compute_state <- function(state, before, condition) {
  (state == 0 && before) || (state == 1 && condition)
}

これをよく見ると、after3 を計算したときの式と全く同一であることがわかる。

すなわち、今回の問題では、状態の計算結果と出力が一致する。

この場合、状態を計算しなくても、直接出力を計算することができる。

compute_after <- function(after, before, condition) {
  (!after && before) || (after && condition)
}

d |>
  mutate(after4 = accumulate2(before, condition, compute_after, .init = FALSE) |> tail(-1))

   before condition after after2 after4
 1 FALSE  FALSE     FALSE FALSE  FALSE 
 2 FALSE  TRUE      FALSE FALSE  FALSE 
 3 TRUE   TRUE      TRUE  TRUE   TRUE  
 4 FALSE  TRUE      TRUE  TRUE   TRUE  
 5 FALSE  TRUE      TRUE  TRUE   TRUE  
 6 FALSE  FALSE     FALSE FALSE  FALSE 
 7 FALSE  TRUE      FALSE FALSE  FALSE 
 8 FALSE  TRUE      FALSE FALSE  FALSE 
 9 TRUE   TRUE      TRUE  TRUE   TRUE  
10 FALSE  TRUE      TRUE  TRUE   TRUE

Enjoy!