Haskell の Array

Haskellのカレンダー | Advent Calendar 2023 - Qiita 3日目の記事です。

Haskell の Array (配列) について書こうと思います。Haskell の Array は索引が型クラスの Ix で抽象化されているため、特に配列の次元を拡張する際に柔軟性がありとても便利です。

そんな便利な Array ですが、もともと Haskell はリスト操作が強力ということもあってか、既存の参考書をみても Array の解説はほんの少しにとどまっているか、解説がないことがほとんどです。

Array が必要になる場面の多くは「リストだと !! によるインデックスアクセスで $O(n)$ になってしまい間に合わない」という場面が多いと思います。しかし Haskell にはインデックスアクセスが $O(1)$ の Vector (vector: Efficient Arrays) もあります。Vector は API が豊富でリストのようにも扱えるし、Stream Fusion などの最適化により何気なく書いた実装が高速化されるなど、非常に強力なライブラリになっています。

競技プログラミングなどの文脈でも Array よりも Vector を好んで使う人は多く、 Array はいまいち日陰感があります。が、前述の通り Array は索引に Ix のインスタンスなら何でも渡すことができるという (強力な) 特長があり、私としては抽象的な関数を定義するにあたって手放せない部品になっています。

というわけで、そんな Array の紹介記事です。

なお、私は普段 AtCoder で競技プログラミングを嗜んでいることもあり、本記事も競技プログラミングの文脈が入ってきますので予めご了承ください。競技プログラミングに興味がない方でも意味のあるよう構成はしたつもりです。

Array の基本的な使い方

配列といえば、他のプログラミング言語でもごく基本となるデータ構造なので、特にマニュアルなど見なくても使えるだろう... と思いたいところですが Haskell の Array には Array, UArray, IArray, MArray, IOArray, IOUArray, STArray, STUArray となんだか種類がいっぱいあり初見殺し感があります。

このあたりの基本的な使い方については以下のスライドによくまとまっていて、おすすめです。

• すごい配列楽しく学ぼう | PPT

実際は

• イミュータブルな配列 (Array, UArray)
• ミュータブルな配列 (IOArray, IOUArray / STArray, STUArray)

の 2 種類に大きくは分類されて、さらにその中が Unbox な値しか扱えないが 正格評価 で速い UArray と、遅延評価ありの Array に分かれています。IArray / MArray はそれぞれ、Unbox とそうでない配列を区別なしで扱えるインタフェースが定義されたモジュールです。

基本は UArray / IOUArray (もしくは STUArray) を選択した上で IArray / MArray の関数で操作して、Unbox でない値を使いたい時に Array / IOArray (もしくは STArray) を使うようにすれば良いです。

Haskell の Array の境界は Ix クラスに抽象化されている

冒頭でも述べた通り Haskell の Array はその索引が Ix クラスで抽象化されています。

配列を使う場合、基本的には配列の境界がどこからどこまでかを、タプルで指定します。
以下の実装の (1, 5) がそれです。

{-# LANGUAGE TypeApplications #-}

import Data.Array.IArray
import Data.Array.Unboxed

main :: IO ()
main = do
  let as = listArray @UArray (1, 5) [10, 20, 30, 40, 50 :: Int]

  print $ as ! 3

つまり 索引の境界を自由に決められる ので、配列を 0 始まりではなく 1 始まりで扱うこともできます。

この境界は Ix 制約を満たしてさえいれば良いので、整数ではなく文字 (Char 型) を索引に使うこともできます。

main :: IO ()
main = do
  let as = listArray @UArray ('a', 'z') [1 .. 26 :: Int]

  print $ as ! 'r' -- 18

accumArray を使うと、文字列のバケットを作ってアクセスするのも簡単です。

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}

import Data.Array.IArray
import Data.Array.Unboxed

main :: IO ()
main = do
  let as = accumArray @UArray (+) 0 ('a', 'z') $ map (,1 :: Int) "abracadabra"

  print $ filter ((> 0) . snd) (assocs as) -- [('a',5),('b',2),('c',1),('d',1),('r',2)]

Ix はある型の連続する部分範囲を整数 (内部的な配列のインデックス)にマッピングするのが主眼の型制約で、ざっくりいうと Eq (同値性) Ord (順序) を判定できれば大体いけます。 Data.Ix のドキュメントを眺めてみると良いでしょう。

さて、ここからが本題です。

多次元配列はどう扱えば良いでしょうか?
多くのプログラミング言語では多次元配列を配列の配列や、リストのリストとして扱いますが Haskell の場合は、以下のように扱います。
構築よりも先にアクセス方法を見てもらう方がわかりやすいでしょう。

-- 2次元配列の2行目3列
print $ as ! (2, 3)

索引を2値のタプルにすることで2次元配列へのアクセスが可能です。

2次元配列の構築を見てみましょう。

main :: IO ()
main = do
  -- 3 x 3 の配列
  -- 1,2,3
  -- 4,5,6
  -- 7,8,9
  let as = listArray @UArray ((1, 1), (3, 3)) [1 .. 9 :: Int]

  print $ as ! (2, 2) -- 5

これで3行3列の2次元配列が表現できています。

listArray に渡している配列の値が、ただのフラットなリストなのが面白いところです。Haskell の配列では、値の方は構造化されておらずフラットで連続する値になっていて、境界である ((1,1), (3,3)) の指定でその構造を指定するわけです。

競技プログラミングをやっていると、入力が二次元のグリッドが与えられてそれに対して何かをするという問題は頻出です。例えば、以下はグリッドを幅優先探索する問題です。

https://atcoder.jp/contests/atc002/tasks/abc007_3

以下のような入力が与えられます。

7 8 -- h行, w列
2 2 -- 開始地点の座標
4 5 -- ゴール地点の座標
########
#......#
#.######
#..#...#
#..##..#
##.....#
########

これを読み取って2次元配列で扱うには、文字列を行ごとに読み取って連結し、先に同じく listArray で二次元配列にしてやれば良いでしょう。
(そのあと幅優先探索をどうするかは、また後ほど)

main :: IO ()
main = do
  [h, w] <- getInts
  s <- getTuple
  t <- getTuple

  css <- replicateM h getLine

  let grid = listArray @UArray ((1, 1), (h, w)) $ concat css

  -- (i, j) でグリッドの座標にアクセスできる
  print $ grid ! (1, 1)
  print $ grid ! s
  print $ grid ! t

{-- 以下はただの入出力 --}

getInts :: IO [Int]
getInts = map (read @Int) . words <$> getLine

getTuple :: IO (Int, Int)
getTuple = do
  [a, b] <- getInts
  return (a, b)

先にみた通り値の方はフラットなので concat で結合するのがポイントです。

二次元がこれで行けるなら三次元も変わらないですね。

main :: IO ()
main = do
  let as = listArray @UArray ((1, 1, 1), (3, 3, 3)) [1 .. 27 :: Int]

  print $ as ! (2, 2, 2) -- 14

境界に指定するのはタプルなので、各次元の型が異なっていても問題ありません。
例えば2次元の座標に加えて、3次元目として探索の進行方向 L, R, U, D を扱いたいとしましょう。以下のように方向を表現する型を宣言し、Ix 制約のインスタンスにすることで配列境界に使うことが可能です。(これを利用した問題もあとで見てみます)

data LRUD = L | R | U | D deriving (Show, Eq, Ord, Ix)

let dist = listArray @UArray ((1, 1, L), (h, w, D)) xs

-- (snip.)

print $ dist ! (2, 2, R)

配列の境界に柔軟性があるということに加えもう一点着目すべきは、 たとえ配列の境界がどんな型であっても配列に対する操作自体は変わらない というところがポイントです。

例えば Data.Ix に定義されている range 関数を使えば Ix で定義された境界の下界から上界までを全探索することができます。
また bounds 関数で配列の境界の下界と上界をタプルで取得できます。

配列の頂点が一次元でも、二次元でも三次元でも、あるいは複数の型が混在していても

print $ map (as !) (range (bounds as))

という関数で、配列の全要素にアクセスすることができます。

このように Haskell の Array は 配列の境界の方が構造の主眼であり、「配列の配列」のようにデータ構造を組み合わせて構造化しているわけではない というところに「なるほど！」感があります。

配列境界の柔軟性をうまく利用する ··· BFS 実装の例

Haskell の配列の配列境界に柔軟性があるのは分かったが、どんなケースで役立つのか例を紹介しましょう。
先にも例を出した幅優先探索の実装例です。

まずは以下のグラフの問題を題材にしてみます。

https://atcoder.jp/contests/tessoku-book/tasks/math_and_algorithm_an

頂点番号 $1 \ldots N$ までの頂点と、それら頂点を結ぶ辺が $M$ 本与えられる。
頂点 $1$ から各頂点への最短経路数を、それぞれ出力する。幅優先探索 (BFS) の典型問題ですね。グラフを構築して、BFS で頂点 $1$ からの最短経路を求めて出力するだけです。(提出した実際のコード)

main :: IO ()
main = do
  [n, m] <- getInts
  uvs <- replicateM m getTuple

  -- グラフの隣接リストを構築
  let g = graph2 (1, n) uvs

	  -- BFS する。dist に配列で各頂点への最短距離が返ってくる
      dist = bfs (g !) (-1) (1, n) [1]

  for_ (elems dist) $ \d ->
    print d

graph2 と bfs がオリジナルの関数です。
それぞれ見てみます。

graph2 は Array を使ってはいますが特に変わった使い方はしていません。u => [v1, v2 ,v3 ... ] という形で、ある頂点から隣接している頂点のリスト (隣接リスト) を accumArray で構築しているだけです。なお頂点同士双方向で結ぶ無向グラフです。

graph2 :: Ix i => (i, i) -> [(i, i)] -> Array i [i]
graph2 b uvs = accumArray (flip (:)) [] b uvs'
  where
    uvs' = concatMap (\uv -> [uv, swap uv]) uvs

次が本丸であるところの BFS の関数です。
BFS のアルゴリズムの詳細や実装は本記事の主眼ではないので気にしなくて良いです。関数のインタフェースに着目してください。

bfs :: Ix v => (v -> [v]) -> Int -> (v, v) -> [v] -> UArray v Int
bfs nextStates initial (l, u_) v0s = runSTUArray $ do
  dist <- newArray (l, u_) initial

  forM_ v0s $ \v0 -> do
    writeArray dist v0 0

  aux (Seq.fromList v0s) dist
  return dist
  where
    aux Empty _ = return ()
    aux (v :<| queue) dist = do
      d <- readArray dist v
      us <- filterM (fmap (== initial) . readArray dist) (nextStates v)

      queue' <-
        foldM
          ( \q u -> do
              writeArray dist u (d + 1)
              return $ q |> u
          )
          queue
          us

      aux queue' dist

この BFS 関数は、グラフの頂点を Int ではなく Ix 制約を満たす型変数 v で定義しています。
そして戻り値は UArray ですが、その境界もこの v を使っています。
BFS の内部では、グラフのどこの頂点に訪問したかを管理して訪問済み頂点への最短距離を随時更新していくわけですが、このデータ構造に STUArray を使っています。この STUArray の境界の型もやはり v を使っています。

グラフの頂点に型変数を使っていて、具体的な型を使っていません。抽象的に頂点を扱っている ということになります。

どういうことでしょうか? もう一つ別の問題、グリッドに対する BFS の問題を解くと、利点が明らかになるので進みます。

途中で解きかけたグリッドBFSの問題を改めて解いてみます。

https://atcoder.jp/contests/atc002/tasks/abc007_3

グリッドの各マスを二次元 $(i, j)$ で表現されたグラフの頂点だとみなして BFS すると、簡単に解けます。(提出コード)

main :: IO ()
main = do
  [h, w] <- getInts
  s <- getTuple
  t <- getTuple

  css <- replicateM h getLine

  let grid = listArray @UArray ((1, 1), (h, w)) $ concat css
      dist = bfs (filter (reachable grid) . around) (-1) (bounds grid) [s]

  print $ dist ! t

around v = map (v `add2`) [(1, 0), (0, 1), (-1, 0), (0, -1)]

reachable grid v
  | (not . inRange (bounds grid)) v = False
  | grid ! v == '#' = False
  | otherwise = True

一つ前の問題に比較すると、

• 一つ前の問題 ··· グラフの頂点は1次元で、ある頂点 v から辿れる次の頂点集合は配列を辿るだけ、つまり g ! v で取得できた
• グリッドになりグラフが二次元になった。ある頂点 v から辿れる頂点集合は v の周囲 4マスのうち進入可能なマスになった。filter (reachable grid) (around v) で取得できる

というところが変わりました。

一方で BFS の実装は変わっていません。

先にもみた通り、 Haskell の Array は境界の構造の方がデータ構造の主眼であって、配列の次元がどうであれ、配列を走査したり更新したりといった計算は次元に依らずに同じです。上記の bfs 関数の実装は内部で STUArray を利用していますが、配列の次元が変わっても計算構造は同じなので変更の必要がありません。一次元配列が二次元配列になったからといって、全探索のため for ループを一つ加える、という変更をする必要はありません。BFS の実装が配列の次元に依存していない、ということです。

もう一つ BFS の例を見てみます。

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

こちらは比較的最近でた BFS の問題で、ポケモンの氷の床 (私は世代的にドラクエですが...) のようなグリッドがあって、壁にぶつかるまで一定方向にしか進むことができず、壁に当たったところで方向転換が可能... というグリッドを舞台にして、到達可能なマスがどこかを計算する問題です。

以下が私の提出コードの抜粋です。(実際のコード)

解説は割愛しますが、data LRUD によって上下左右 4方向のデータ型を定義し、BFS は方向まで加味した三次元の状態遷移を高階関数で渡すことによって問題を解いています。ここで使っている bfs 関数もまた、先の実装に同じです。bfs 関数を一切変更することなく、状態遷移の関数を問題設定に合わせて定義することで問題を解くことができます。

data LRUD = L | R | U | D deriving (Show, Eq, Ord, Ix)

around grid (x, y, d) = do
  let u@(x', y', _) = case d of
        L -> (x - 1, y, L)
        R -> (x + 1, y, R)
        U -> (x, y - 1, U)
        D -> (x, y + 1, D)

  if grid ! (x', y') == '#'
    then [(x, y, d') | d' <- [L, R, U, D]]
    else [u]

reachable grid (x, y, _)
  | (not . inRange (bounds grid)) v = False
  | grid ! v == '#' = False
  | otherwise = True
  where
    v = (x, y)

main :: IO ()
main = do
  [h, w] <- getInts
  grid <- getCharsGrid ((1, 1), (h, w))

  let v0s = [(2, 2, d) | d <- [L, R, U, D]]
      dist = bfs (filter (reachable grid) . around grid) (-1) ((1, 1, L), (h, w, D)) v0s

      s =
        foldl'
          (\acc (x, y, _) -> Set.insert (x, y) acc)
          Set.empty
          (map fst . filterOnSnd (>= 0) $ assocs dist)

  print $ Set.size s

繰り返しになりますが、BFS の実装が配列の次元に依存していないことにより、探索の計算構造の抽象度が上がっており、すべての探索問題に同じ実装で対応できていることがわかります。配列の次元が抽象化されているのが、こういうところで役に立ちます。

$N$ 次元ナップザック問題

配列などインデックスアクセス可能なデータ構造に計算過程をメモ化しながら網羅的な計算をやっていくアルゴリズムの代表例といえば、やはり動的計画法 (DP) です。
DP では更新可能な配列を用意し、 for ループを配列の次元に合わせて回しながら配列を手続き的に更新していくというのが一般的かと思います。

一方、DP の中でもナップザック問題のように状態空間全体を計算の状態遷移に合わせて毎回更新していくような構造は、畳み込みで計算することもできます。私はこの畳み込みで解ける DP に対しては以下のような関数を用意して対応しています。

ここでも実装の詳細は気にしなくて良いです。インタフェースを見ると Ix の型変数が使われていて、内部実装も accumArray や bounds や inRange といった配列操作の関数が出てはくるが二重、三重の for ループみたいなものはどこにもない、というところだけ雰囲気を見てみてください。

-- 畳み込みDP
-- 時間的遷移時に、状態空間が前の状態を引き継ぐ (accum)
-- ex) accumDP @UArray f max minBound (0, wx) [(0, 0)] wvs
accumDP ::
  ( IArray a e,
    Ix v,
    Eq e,
    Show e,
    Show v,
    Show (a v e),
    Foldable t
  ) =>
  ((v, e) -> x -> [(v, e')]) -> -- 状態遷移関数 f v / x をみて v の次の遷移可能性を返す
  (e -> e' -> e) -> -- 緩和の二項演算
  e -> -- 初期値 (0, minBound, maxBound, False など)
  (v, v) -> -- 状態空間の下界、上界
  [(v, e')] -> -- 開始時点の状態
  t x -> -- 入力 (時間遷移)
  a v e -- Array or UArray
accumDP f op initial (l, u) v0s xs = do
  let dp = accumArray op initial (l, u) v0s
  foldl' transition dp xs
  where
    transition dp x =
      accum op dp $
        filter (inRange (bounds dp) . fst) $
          concatMap (`f` x) (assocs dp)
      where
        !_ = dbg ("dp", filter ((/= initial) . snd) $ assocs dp)

典型的なナップザック問題は以下になります。

https://atcoder.jp/contests/tessoku-book/tasks/tessoku_book_s

手続き的に解くと二次元DP ですが、上記の関数では状態空間と同じサイズの一次元の配列を一つ用意してそれを遷移に合わせて更新していくのみになります。

main :: IO ()
main = do
  [n, wx] <- getInts
  wvs <- replicateM n getTuple

  let dp = accumDP @UArray f max minBound (0, wx) [(0, 0)] wvs
        where
          f (w, v) (wi, vi)
            | v == minBound = []
            | otherwise = [(w + wi, v + vi)]

  print $ maximum (elems dp)

これで問題が解けます。
DPの初期設定として状態空間の下界と上界を (0, wx) で指定し、初期値や緩和時の二項演算、初期状態なども合わせて設定します。
その上で、状態遷移の関数 $f$ を定義する。これだけです。

この簡単なナップザック問題を解くだけならわざわざ抽象的な関数を用意する必要はないわけですが、問題が複雑になってくるとその威力がよくわかると思います。

先のナップザック問題は状態空間の次元が一次元で済んでいましたが、以下のナップザック問題の亜種である問題は最大 $5$ つのパタメータの組み合わせ全てを探索する必要がある DP の問題で、力技でいくと 5次元あるいは6次元 DP を解くことになります。

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

5次元 DP を記述するのは結構大変なので、5 つのパラメータの冪集合を $N$ 進数で表現するテクニックを使うのが典型例っぽいですが、拙作のフレームワークであれば例によって次元が上がっても計算構造自体は変わらないので、先のシンプルなナップザック問題と変わらず解くことができます。accumDP の実装には一切手を加えていません。

main :: IO ()
main = do
  [n, k, p] <- getInts
  xs <- replicateM n getInts

  let xs' = map (\x -> x ++ replicate (5 - k) 0) xs

  let dp = accumDP @UArray f min (maxBound :: Int) ((0, 0, 0, 0, 0), (p, p, p, p, p)) [((0, 0, 0, 0, 0), 0)] xs'
        where
          f ((p1, p2, p3, p4, p5), v) [c1, a1, a2, a3, a4, a5]
            | v == maxBound = []
            | otherwise = [((min p (p1 + a1), min p (p2 + a2), min p (p3 + a3), min p (p4 + a4), min p (p5 + a5)), v + c1)]
          f _ _ = error "!?"

  let x =
        if
            | k == 1 -> dp ! (p, 0, 0, 0, 0)
            | k == 2 -> dp ! (p, p, 0, 0, 0)
            | k == 3 -> dp ! (p, p, p, 0, 0)
            | k == 4 -> dp ! (p, p, p, p, 0)
            | k == 5 -> dp ! (p, p, p, p, p)
            | otherwise -> error "over 5"

  print $ if x == maxBound then -1 else x

配列の境界の次元が 5 次元になってもどこにも 5 重ループ的なものが出てこないのが良いですね。
実際、上記のコードで本番中にこの問題を通すことができました。

まとめ

Haskell の配列は、一見すると API も少なく素朴です。
Vector はインデックスアクセスもできて、速くて、いろんなことができるのに Array には API がほんの 10数個程度しかありません。

しかし、配列の境界が Ix クラスで抽象化されていて、Ix も素朴ながら僅か数個の API で、配列アクセスを容易にする機構を持っています。
たかだか十数個程度の関数で構成されている Array と Ix をうまく利用すると、ここまでにみた通り、アルゴリズムに強力な抽象化を施すことができます。シンプルでありながら強力、という点では Haskell のコンセプトをよく体現した設計だなあと思います。

ぜひ Array を道具箱の特等席に置くことも検討してみてください。