永続データプログラミングと競技プログラミング 〜 Haskell でがんばる競プロ

https://x.com/naoya_ito/status/1864254832074498335

この記事は Haskell - Qiita Advent Calendar 2024 - Qiita の 5日目の記事です。

純粋関数型言語の Haskell では、値は基本的に不変です。リスト、Set、Map など基本的なデータ構造も不変データ構造として提供されています。

不変なデータ構造は変更をしても、変更前の値が残ります。結果、データを変更したあとも以前のデータを参照することができるという特性が得られます。この特性を指して、不変なデータのことを「永続データ」と呼び、永続データを駆使して問題を解くことを「永続データプログラミング」と呼ぶことがあります。

Haskell で関数型プログラミングをすると、自然と永続データプログラミングを実践することになります。

永続データプログラミングを支えるのが「永続データ構造」の存在です。

不変なデータで変更を表現するには、元のデータをコピーして、それを元に変更後のデータを作るという方法が考えられます。このときデータ構造全体をナイーブにコピーしてしまうと、抱えているデータのサイズ分だけコピーが行われて非効率です。

データ構造全体をコピーするのではなく「変更されるノードとそのノードへ直接的・間接的に参照を持つノードだけをコピーする」ことによって計算量を抑え、不変でありながらも効率的なデータ更新が可能になるよう実装されているのが、永続データ構造です。

Haskell のリスト、Set、Map など不変なデータ構造はいずれも永続データ構造であると言えます。不変なデータ構造を使って実装していくと、それは必然的に永続データプログラミングになる、と言えます。

永続データプログラミングと永続データ構造についてのあらましは、先日 永続データプログラミングと永続データ構造 - 一休.com Developers Blog に解説しましたので、参考にしてください。

https://user-first.ikyu.co.jp/entry/2024/12/03/091133

Haskell で競技プログラミング

さて、そんな Haskell で競技プログラミングをする話です。

ご存知の通り、Haskell では変更可能なデータ構造を用いた命令型 (手続き的) プログラミングをすることも可能ではあります。

が、わざわざ Haskell を使って競技プログラミングを嗜むのであれば、私個人としては、なるべく関数型プログラミングで問題を解いていくことを目標としたいところです。イミュータブルすなわち永続データプログラミングのスタイルで、競技プログラミングの問題を解く!

競技プログラミングでは量的に大きなデータを扱うので、イミュータブルにやっていくなら永続データ構造が必須です。

命令型でやればメモリ使用量的にも実行速度的にも効率的であることがわかっていながら、あえて関数型プログラミングで競プロに挑むのは果たしてハンデなのでしょうか? いいえ、永続データプログラミングを意識的に実践することにより、それがアドバンテージになる場面というのは多々あります。

本稿では Haskell、関数型プログラミング、永続データプログラミングで競技プログラミングに挑むとはどんな感じなのか、私なりの考え方を紹介していきたいと思います。

永続データを意識して問題を解く 〜 命令型プログラミングとの対比

簡単な問題を取り上げます。

https://atcoder.jp/contests/abc081/tasks/abc081_b

$N$ 個 ($1 \leq N \leq 200$) の正の整数 $A_1, \ldots, A_N$ がある。
整数がすべて偶数なら全部 $2$ で割ったものに置き換える。最大で何回この操作を行うことができるか?

8 12 40

という数列があったら、2 で割るたび

4 6 20
2 3 10

と遷移したところで奇数が出る。そこで遷移は終わり。
この例では最大の操作回数は $2$ 回です。

$N$ が小さいので、問題文のとおりシミュレーションし、要素すべてが偶数ならすべて $2$ で割るということを繰り返し、奇数が出るまでの繰り返し回数を数えるというナイーブな実装で解けそうです。

これを命令的 (手続き的) に実装する場合、たとえば以下のようになるでしょう。Python で、命令的な記述を ChatGPT に書かせました。

n = int(input())
as = list(map(int, input().split()))

count = 0

while all(a % 2 == 0 for a in as):
    for i in range(n):
        as[i] //= 2
    count += 1

print(count)

カウンターを用意して、条件を満たしている間 ··· つまりすべて偶数である間 for ループで繰り返し要素を $2$ で割っていき、カウントアップするという実装です。

さて、私が Haskell で実装する場合はすこし様子が異なります。

main :: IO ()
main = do
  _ <- getInt
  as <- getInts

  let res = iterate (map (`div` 2)) as

      i = fromJust $ findIndex (any odd) res

  print i

とするか、もしくは

main :: IO ()
main = do
  n <- getInt
  as <- getInts

  let res = takeWhile (all even) $ iterate (map (`div` 2)) as

  print $ length res

のように実装します。

ポイントになるのは、リストの要素を $2$ で割る写像を繰り返す計算と、奇数が出てくるタイミングを見つける計算を分離している点です。

最初の命令的な実装では、操作のたび配列の値を破壊的に変更しています。

一方、私の Haskell の実装は、リストの元の値を直接変更するのではなく、写像 (map) によってリストから別のリストを作っています。iterate で写像を繰り返すことでリストの変更履歴が作られます。

このリストの履歴を生成してから、改めて履歴をみて奇数を含むのが何番目か findIndexで検索する、あるいは偶数の間写像を繰り返してできあがった履歴の長さを数える、ということで答えとしています。

命令的な実装ではその場その場で元の配列を書き換えるため、データ構造を変更してしまったら、変更前の状態にはアクセスできません。短命データを使った命令型プログラミングです。Haskell の方は、繰り返しリストの写像を行う過程がすべて残るという永続データになっている。永続データプログラミング、と言えます。

「何かをしながら、何かをする」or「何かをしたあと、何かをする」

もう少し両者の違いを見てみます。

命令的な実装では、ループを回し配列を書き換えながら、カウンターをカウントアップしています。

while all(a % 2 == 0 for a in as):
    for i in range(n):
        as[i] //= 2
    count += 1

Haskell の方は、

  let res = iterate (map (`div` 2)) as
      i = fromJust $ findIndex (any odd) res

写像の繰り返しと、操作回数を求める探索の計算を分離しています。

短命データは書き替えを行うとそれ以前の状態にはアクセスできません。処理効率は良いですが、ある時点のデータ構造にはその時点にしかアクセスできないので、その時々のデータ構造の状態に応じた別の計算が必要な場合、その二つの計算を切り離すことが難しい。必然的に「何かをしながら、何かをする」という実装が多くなります。「何かをするついでに」とも言えるでしょう。今回の場合は「配列を書き換えながら、カウンターを更新する」となります。

一方の永続データプログラミングでは変更前のデータにアクセスすることが可能なので、データ構造の変更 (状態遷移) と後続の処理を分離することが可能です。これにより「何かをしたあと、何かをする」という実装が可能です。今回の場合は「リストを繰り返し写像し終えてから、数える」となります。永続データが計算と計算の分離に一役買っているわけです。

もうひとつ、似たような問題を見てみましょう。

https://atcoder.jp/contests/abc311/tasks/abc311_a

A, B, C からなる文字列 $S$ が与えられる。$S$ には A, B, C しか登場しない。
この $S$ を左から $1$ 文字ずつ見ていったときに、はじめて A, B, C すべてが $1$ 回以上出現するのは、左から何文字目まで見たときか?

たとえば

ACABB

は、$4$ 文字目の B ではじめて ABC が出たことになるので、答えは $4$ です。

命令的に実装したものは、例えば以下のようになるでしょう。こちらも ChatGPT に書かせました。
set を利用して、$3$ 種類の文字が出てくるタイミングまでループを続けて、$3$ つ出てきたところで答えを出力して break しています。

n = int(input())
s = input()

seen = set()

for i in range(n):
    seen.add(s[i])
    if len(seen) == 3:
        print(i + 1)
        break

先の問題に同じく、命令型の実装は条件を満たすまでループを回し虎視眈々とそのタイミングを伺い、時がきたところで答えを出力しています。「set を更新しながら、3 文字揃うのを待つ」プログラムです。

一方の私の Haskell の実装は以下です。

main :: IO ()
main = do
  _ <- getInt
  s <- getLine

  let ss = scanl (flip Set.insert) Set.empty s
      i = fromJust $ elemIndex (Set.fromList ['A', 'B', 'C']) ss

  print i

こちらも Set を使いましたが、Haskell の Set は永続データ構造であり、変更に伴う状態遷移の過程をすべて残すことができます。これにより文字列 $S$ をみながら Set を変更していくという計算と、A, B, C 三文字が揃うタイミングがいつかを見つける計算を分離できます。「$S$ に合わせて Set を変更するのを終えてから、3 文字揃うタイミングを検索する」実装になっています。

(命令型の実装は 3 文字揃ったところで break しているので効率は良いでしょう。なお、計算量的には同じ $O(n)$ です。1 度のループで処理できている命令型の実装と、2 度の再帰が走っていそうに見える Haskell ですが遅延評価があるのでその点は単純には比較できません。)

このように命令型の実装は「何かをしながら、何かをする」実装になりやすい一方、関数型の実装は「何かをしたあと、何かをする」実装にしやすい、という特徴があります。

Haskell でプログラムを実装するときの一つの目標として、計算をなるべく独立した単位に分離し分割統治することがあると思います。その分離された小さな計算は多くの場合、写像か、畳み込みです。「写像してから、畳み込む」つまりmap / reduceですが、これが関数型プログラミングの典型的なメンタルモデルであり、競技プログラミングの実装も map / reduce に落とし込んでいくのが実装の基本になります。

おそらく多くの関数型プログラマーと共有できる感覚ではないかと思います。

「永続データ構造」が活きる

先の2問は対象としているデータが小さいものでした。ナイーブにデータ構造を写像しても特段問題にはならないでしょう。

では、もう少しデータが大きくなってきたときにも同様の考え方で実装することはできるのでしょうか? 結論、できます。

次の問題を解いてみます。先の 2 問に比較するとデータサイズも大きく、難しい問題です。

https://atcoder.jp/contests/abc306/tasks/abc306_e

長さ $N$ ($1 \leq N \leq 5 \times 10^5$) の数列 $A$ があり、最初は全ての項が $0$ である。
数列 $A$ の $A_{X_i}$ を $Y_i$ に変更しろ、というクエリが $Q$ ($1 \leq Q \leq 5 \times 10^5$) 回次々と指示される。数列を変更するたび、数列に含まれる要素のうち大きなもの top $K$ 個の和をとって出力せよ。

1 5
2 1
3 3
4 2
2 10

とクエリが与えられます。最初に $A_1 =5$ にして、次に $A_2 = 1$ にして... という指示ですね。
$N = 4$ だったとすると、その都度数列が

$(5, 0, 0, 0)$
$(5, 1, 0, 0 )$
$(5, 1, 3, 0)$
$(5, 1, 3, 2)$

と更新されていくので、$K = 2$ の場合

5
6
8
8

が、各タイミングでの大きな数 top $K$ の和ということになります。

数列を更新しながら、そのときどきの top K の和を求める、ということでいかにも「データ構造を変更しながら、その時点の状態に応じた値を得る」ことをやる問題です。

この問題は、中核で使うデータ構造を集合として考えるとよいです。$A_{X_i}$ の値が書き換えられるたび、集合から以前の値が削除されて、新しい値が集合に追加されると考えられます。含まれる値の top $K$ を常に保持できるような集合の実装を持っていれば、それを使って解くことができます。

私は BalancedMultiSet という名前で top $K$ を保持する集合のライブラリを自作しています。

BalancedMultiSet の実装は、内部的には Data.IntMap を二つ使ったイミュータブルな実装になっているため、BalancedMultiSet もまた永続データ構造です。集合に値が追加されたり、削除されるたびに二つの IntMap を探索して必要に応じてそれぞれの最大 / 最小値を交換することにより、一方に top K の値が集まるようにします。本記事の趣旨とこの集合ライブラリの実装はあまり関係がないのでより詳細は割愛しますが、実装に興味がある方は提出済みのコード (提出 #60431339 - トヨタ自動車プログラミングコンテスト2023#3(AtCoder Beginner Contest 306)) をみてください。

main 関数の実装は以下のようになります。

main :: IO ()
main = do
  [n, k, q] <- getInts
  qs <- replicateM q getTuple

  as <- newArray @IOUArray (1, n) (0 :: Int)

  -- 空の BalancedMultiSet (集合) を用意
  let s0 = emptyBMS k GT

  ss <- scanForM s0 qs $ \s (xi, yi) -> do
    prev <- readArray as xi
    writeArray as xi yi

    -- 集合から Ai の以前の値を削除し、新しい Ai を追加する
    return $ (insertBMS yi . deleteBMS prev) s

  -- 集合の状態遷移の履歴をすべて top K の和に変換する
  let res = [leftSum s | s <- tail ss]

  for_ res print

{-- BalancedMultiSet --}

data BalancedMultiSet = BalancedMultiSet
  { topKValue :: Int,
    topKOrder :: Ordering,
    left :: IntMultiSet,
    right :: IntMultiSet,
    leftSum :: Int
  }
  deriving (Show)

emptyBMS :: Int -> Ordering -> BalancedMultiSet
emptyBMS k order =
  BalancedMultiSet
    { topKValue = k,
      topKOrder = order,
      left = emptyMS,
      right = emptyMS,
      leftSum = 0
    }

-- BalancedMultiSet の実装詳細は略 --

これまでにみてきた問題同様の方針、つまり「集合の状態遷移をすべて済ませたあと、その履歴から答えを得る」ようにしています。

具体的には

ss <- scanForM s0 qs $ \s (xi, yi) -> do
    -- snip. --
    return $ (insertBMS yi . deleteBMS prev) s

ここで繰り返し、集合から古い値を削除し新しい値を追加するという状態遷移を行っています。

この問題で最終的に得たい値は top $K$ の値の和ですが、その和の値を得ることは同時にはしません。集合の状態を遷移させることに徹して、それを終えるところまでやりきっています。変数 ss には、集合が空である初期状態からクエリを一つ処理するたび集合の状態が遷移する過程すべてがリストとして束縛されています。

これでクエリを処理したその時々の集合の状態は得られるので、

    let res = [leftSum s | s <- tail ss]

と、改めてデータ構造それぞれの状態を top $K$ 個の値の総和に写像し、答えを得ています。

先の問題に同じく「何かをし終えてから、何かをする」と切り離して実装できています。

さて、この問題ではクエリが最大で $5 \times 10^5$ 回も発行されます。つまり、集合の状態遷移は $5 \times 10^5$ 回行われて、さらに集合には値がクエリと同じ回数だけ挿入されるので集合のサイズも最大で $5 \times 10^5$ 個の値を抱えることになります。 $5 \times 10^5$ 個のデータを抱える大きな集合を、$5 \times 10^5$ 回状態遷移させた履歴を永続データとして持つ... なかなか大きなデータになりそうですが、問題ないのでしょうか?

ご想像のとおり問題ないからこそ、TLE もしくは MLE にならずに、上記の実装で通っているわけです。BalancedMultiSet の実装で利用している Data.IntMap は永続データ構造です。

永続データ構造は冒頭で述べたとおりデータの変更にあたり直接・間接的に変更に関係のあるノードのみを変更して、それ以外は変更前後でノードを共有します。従って $5 \times 10^5$ 個の集合 (への参照) があったとしても、その大部分の領域は共有されています。

永続データ構造であれば、たとえ $5 \times 10^5$ オーダーの変更であっても、永続データプログラミングをしても問題がおきません。

雑に言えば、Haskell では Set、Map などのデータ構造を写像により大量に状態遷移させても大丈夫だし、その考えに基づいて実装をすることができる、ということです。(もちろん、その時々の計算量を考える必要はあります)

よってこのケースでも、最初の2問同様に「何かをしおえてから、何かをする」「写像してから、畳み込む」というようにやりたいことを分離して実装することができます。

各々分離して実装できるということは、一つのコードブロックをより少ない責務にフォーカスさせて実装することができる、つまり、より認知負荷低く実装することができることを意味します。また読むときも、より認知負荷低く読むことができると思います。永続データ構造もまた、計算と計算を分離するのに役立っているわけです。

そして、処理を分離できるということはプログラムの構造をフラットに保ちやすいとも言えます。

命令型プログラミングの場合、データ構造のその時々の状態に依存した処理が連なって分離が難しいため、二重ループ、三重ループと入れ子になりがちですが、Haskell の場合、永続データ (それに加えてモナド) があるため、処理の入れ子を生じさせずフラットに実装することができます。

永続データプログラミングで ABC322 D - Polyomino を解く

ここまで見てきた方針を活かして、かの悪名高き? ABC322 D - Polyomino を解いてみましょう。

https://atcoder.jp/contests/abc322/tasks/abc322_d

この問題は $4 \times 4$ サイズのパズルのピースが $3$ 種類与えられる中、それを組み合わせて $16 \times 16$ の正方形を作ることができるか? を判定する問題です。
パズルのピースのことを「ポリオミノ」と呼ぶそうです。ポリオミノは上下左右に動かしたり、回転したりはしてよいですが、ポリオミノの一部分を重ねて配置することはできません。

$3$ つの図形それぞれの上下左右 $4$ マスの移動、$0, 90, 180, 270$ 度回転すべてを試す、つまり全探索で解ける問題なのですが、手続き的に何気なく for ループを書き始めると入れ子が何重にもなって実装が大変です。実際、グーグルで検索してみてみると、入れ子の○重ループの実装で解いている例をたくさん見かけます。

「D問題のわりに重実装問題」として悪名高いのがこの問題です。

しかし、ここまで見てきたとおり Haskell で「何かをし終えてから、何かをする」という分割統治の考え方に従って実装していくと、シンプルに実装することができます。
ポイントになるのは

1. 二次元のグリッドを (配列ではなく) 集合 (Set) で表現する
2. 3つのポリオミノそれぞれを 4回転させて、上下左右 4方向に移動させた Set をすべて先に用意する。つまり、$3$ 種類、回転 $4$ 回、$4$ 方向に $4$ マス並行移動の $3 \times 4 \times 4\times 4 = 192$ 個の形すべてを写像で全部先に作ってしまう。
3. できあがった全パーツから、3つを選び Union し、正方形判定をする

の、特に 2 の点です。

isSquare x y z
  | not (Set.disjoint x y && Set.disjoint x z && Set.disjoint y z) = False
  | Set.size p /= 16 = False
  | otherwise = all (\v -> Set.member v p) $ range ((1, 1), (4, 4))
  where
    p = foldl1 Set.union [x, y, z]

main :: IO ()
main = do
  gridA <- getCharGrid ((1, 1), (4, 4))
  gridB <- getCharGrid ((1, 1), (4, 4))
  gridC <- getCharGrid ((1, 1), (4, 4))

  let a = fromArrayGS (== '#') gridA
      b = fromArrayGS (== '#') gridB
      c = fromArrayGS (== '#') gridC

      ds = map (\[dx, dy] -> (dx, dy)) $ sequence [[0 .. 3], [0 .. 3]]

      -- 回転して上下左右に並行移動した部品をすべて先に作る
      as = [a'' | a' <- take 4 $ iterate rotateRightGS a, d <- ds, let a'' = Set.map (+ d) $ normalizeGS a']
      bs = [b'' | b' <- take 4 $ iterate rotateRightGS b, d <- ds, let b'' = Set.map (+ d) $ normalizeGS b']
      cs = [c'' | c' <- take 4 $ iterate rotateRightGS c, d <- ds, let c'' = Set.map (+ d) $ normalizeGS c']

      -- うち、3つを組み合わせる全探索。正方形判定をする
      res = [isSquare x y z | x <- as, y <- bs, z <- cs]

  printYn $ or res

なお fromArrayGSはグリッドの配列を Set に写像する関数、rotateRightGS は座標の回転写像、normalizeGS は図形を左上マス開始にする正規化の関数です。いずれもたいしたことはやっていません。実装の詳細に興味のある方は 提出 #59487824 - AtCoder Beginner Contest 322 を見てください。

この問題のデータは小さいので永続データ構造の効率性が必要な問題ではありません。一方、永続データや永続データ構造を意識できていると「Set を先に全パターン分作ってしまえばいい」「それをしてから、改めて全探索をすればいい」という 🧠 で考えられます。

図形を命令的データ構造で表現していると、回転したり上下左右移動するたびに以前の形は失われてしまうため、その都度、正方形判定をする必要があり、結果何重もの入れ子のループになってしまうところ、永続データであれば、それを避けてフラットに実装することができるわけです。

ついでに ABC307 C - Ideal Sheet も解く

図形の並行移動といえば C 問題にも関わらず水色難易度だった以下の問題があります。

https://atcoder.jp/contests/abc307/tasks/abc307_c

問題の詳細は割愛しますが、こちらも先の問題同様の考え方で解くと、実装も含め非常にシンプルです。

main :: IO ()
main = do
  [ha, wa] <- getInts
  gridA <- getCharGrid ((1, 1), (ha, wa))
  [hb, wb] <- getInts
  gridB <- getCharGrid ((1, 1), (hb, wb))
  [hx, wx] <- getInts
  gridX <- getCharGrid ((1, 1), (hx, wx))

  let a = fromArrayGS (== '#') gridA
      b = fromArrayGS (== '#') gridB
      x = fromArrayGS (== '#') gridX

      res =
        [ isShiftedGS c x
          | [dx, dy] <- sequence [[-100 .. 100], [-100 .. 100]],
            let a' = Set.map (+ (dx, dy)) a
                c = Set.union a' b
        ]

  printYn $ or res

「何かをしながら、何かをする」のではなく「何かをしたあと、なにかをする」。永続データや永続データプログラミングを意識した実装が、実装をフラットかつシンプルに保つきっかけになっていることが伝わるでしょうか?

まとめ

Haskell、というよりは関数型プログラミング、不変データ、永続データプログラミングで競技プログラミングを解くとはどういうことかを、見てきました。

命令型プログラミングの場合に比較して、永続データプログラミングを意識することで、個々の計算を独立したより細かい関数に分離することができました。問題によっては、それによって実装の認知負荷 🧠 を大きく下げることができます。決められた時間内に正確な実装を求められる競技プログラミングにおいては、実装の認知負荷を下げ、シンプルに留められることは大きな武器になります。シンプルに実装できるということは、解法の考察も複雑にならずに済む、ということでもあります。

そしてこの計算と計算の分離を可能にする一つのプログラミング要素が永続データであり、計算量を犠牲にせず永続データプログラミングを可能とするため、永続データ構造がそれを支えているのだろう、というのが私なりの考察です。

継続して競技プログラミングの問題を解いていると、関数型プログラミングの強力さに気づく機会が多くあります。関数型プログラミングで競プロをやるというのは、決して枷ではなく、場合によっては命令型プログラミングよりも (遙かに) 有利に問題を解けることもある、と実感しています。

···ただし、その域に達するまでが大変だけどな!