これは何?

本アドベントカレンダーの6日目の記事です。

本記事では、Haskellのリストから要素を取り出す際に$O(1)$になるパターンと$O(n)$になってしまうパターンについて、なぜこの差が生まれるのかとその対処法について説明します。

Haskellのリストは単方向連結リストである

In Haskell, lists are one of the most important data types as they are often used analogous to loops in imperative programming languages. These lists are singly linked, which makes them unsuited for operations that require O(1) access. Instead, they are intended to be traversed.¹

Haskellのリストは単方向連結リストになっている。
そのため、要素が前から順にしか値がとりだせない構造になっている。

[1, 2, 3]のようなリストは内部では`1 : (2 : 3 : [])のようにして保存されている。

ghci> :t [1,2,3]
[1,2,3] :: Num a => [a]
ghci> :t 1 : (2 : 3 : [])
1 : (2 : 3 : []) :: Num a => [a]
ghci> show $ 1 : (2 : 3 : [])
"[1,2,3]"

先頭の要素を取り出す処理はO(1)になる

headやtailは$O(1)$で実行できる。
これは先頭の要素を取り出すのみで実現できる。

head' :: [a] -> a
head' (x : xs) = x

tail' :: [a] -> [a]
tail' (_ : xs) = xs -- 先頭とそれ以外にわけてそれ以外を返す

main :: IO ()
main = do
  let xs = [1, 2, 3] -- 内部的には 1 : (2 : 3 : [])
  -- 先頭文字列を取得
  print $ head xs -- O(1)
  print $ head' xs -- 説明用に再定義した関数を使用
  -- 先頭以外を取得
  print $ tail xs -- O(1)
  print $ tail' xs -- 説明用に再定義した関数を使用

先頭の要素以外を取り出す必要がある処理は遅くなる

逆にinitやlastのような処理は$O(n)$になるため遅くなる。

init' :: [a] -> [a]
init' [x] = []
init' (x : xs) = x : init' xs

last' :: [a] -> a
last' [x] = x
last' (_ : xs) = last xs

main :: IO ()
main = do
  let xs = [1, 2, 3] -- 内部的には 1 : (2 : 3 : [])
  -- 末尾以外を取得
  print $ init xs -- O(n) init'参照
  print $ init' xs -- 説明用に再定義した関数を使用
  -- 末尾を取得
  print $ last xs -- O(n) -- lastは再帰しないと最後の要素までたどり着かないので遅い(last'参照)
  print $ last' xs -- 説明用に再定義した関数を使用

末尾要素に対するアクセスを高速にする方法

最初に配列を`reverse`で逆さまにしてしまう

この問題は末尾からマッチする文字列があるかを探す必要があるのだが、以下のように素直に末尾からマッチするものを探すと、1回の調査で$O(n)$になるため、全体では$O(n^2)$になってしまい、LTEになってしまう。

-- TLEになる例
import Control.Arrow ((>>>))
import Data.List (isSuffixOf)

-- isSuffixOfを使って後ろからwordListにマッチするものがないか探す。
solve :: String -> Bool
solve "" = True -- 再帰で全ての文字列を消せるならS=T
solve s =
  case filter (`isSuffixOf` s) wordList of -- TODO: O(n)なので改善が必要
    [] -> False -- 末尾がマッチしない場合
    (word : _) -> solve (take (length s - length word) s) -- 末尾からword分だけ除去して再帰
    -- TODO: takeとlengthがO(n)なので再帰が多い場合にはコストがかさむ。wordListの最短文字列が5なのでn + n-5 + n-10 ... = n^2/2 = O(n^2)
  where
    wordList = ["dreamer", "dream", "eraser", "erase"]

main :: IO ()
main =
  interact $ \inputs ->
    let s = init inputs -- initで改行を外す
     in if solve s then "YES\n" else "NO\n"

先に全体をreverseで反転しておくことで先頭から要素をマッチさせることができ、1回の試行が$O(1)$になるため、全体としては$O(n)$になり、LTEを回避できる。

import Data.List (isPrefixOf)

-- NOTE: wordListは大きい順にしておく必要がある。eraserがある時にeraseでマッチすると壊れるから。go (reverse s)をしている関係上各単語の語順を逆にしている。
wordList :: [String]
wordList = map reverse ["dreamer", "dream", "eraser", "erase"]

-- 後ろからwordListにマッチするものがないか探す(一部高速化)
-- dreameraserの場合、前方でマッチするとdreamerがdreamより先に試行されるため、dreamerがマッチするとaserが残ってしまう。
-- 逆にdreamを先に試行させるとdreamerの場合erが残ってしまう。
-- このことから、後方からマッチするものがないものを探すかつ、長い文字列から調査する形をとっている。
solve :: String -> Bool
solve s = go (reverse s)
  where
    go "" = True -- 再帰で全ての文字列を消せるならS=T
    go sReverse =
      case filter (`isPrefixOf` sReverse) wordList of -- NOTE: isSuffixOfを使うと毎回O(n)になるので最適化
        [] -> False -- 末尾がマッチしない場合
        (word : _) -> go (drop (length word) sReverse) -- 末尾からword分だけ除去して再帰(O(n))に最適化

main :: IO ()
main =
  interact $ \inputs ->
    let s = init inputs -- initで改行を外す
     in if solve s then "YES\n" else "NO\n"

Data.Vectorのような配列を使う

Data.Vector²を使うとHaskellで配列を使用できる。

これはいわゆるC言語の配列のように連続したメモリ領域に要素を並べて格納する方式なので、インデックスiに対する要素は$$先頭アドレスの位置 + i * 要素のサイズ$$でアクセスできるため、全ての要素に対して$O(1)$でアクセスできる。

import Data.Vector qualified as V

main :: IO ()
main = do
  let xs = [1, 2, 3] -- 内部的には 1 : (2 : 3 : [])
  let v = V.fromList xs -- 配列なのでどの要素にアクセスしてもO(1)
  -- 末尾以外を取得
  print $ V.slice 0 (V.length v - 1) v -- O(1) NOTE: Data.Vectorは配列ベースなので内部で配列の長さを保持しているため、lengthがO(1)で取得できる。また、リストのようにスライスする際には元の配列を参照するので
  -- 末尾を取得
  print $ V.last v -- O(1)

Reference