連長圧縮したリスト上の分割(Haskell)

はじめに

AtCoder Beginner Contest 413 C Large Queueを見ていて、連長圧縮したリスト上でsplitAtできるといいね、と思ったので実装してみました。汎用性がありそうなので、誰かがすでに実装してそうですが、車輪の再発明も楽しいものなのでよしとしましょう。

連長圧縮

https://ja.wikipedia.org/wiki/連長圧縮

連長圧縮 toRLE

module RLE where

import Data.List

type RLE a = [(a, Int)]

{- |
>>> let xs = "aaabcc"
>>> toRLE xs
[('a',3),('b',1),('c',2)]
-}
toRLE :: Eq a => [a] -> RLE a
toRLE = unfoldr psi where
    psi = \ case
        x:xs -> case spanCount (x ==) xs of
            (m,ys) -> Just ((x, succ m), ys)
        _    -> Nothing

spanCount :: (a -> Bool) -> [a] -> (Int, [a])
spanCount p = \ case
    x:xs | p x -> case spanCount p xs of
        (m,ys)     -> (succ m, ys)
    xs         -> (0,xs)

連長圧縮からの復元 fromRLE

{- |
>>> let xs = "aaabcc"
>>> fromRLE (toRLE xs) == xs
True
-}
fromRLE :: RLE a -> [a]
fromRLE = (uncurry (flip replicate) =<<)

連長圧縮したリスト上の分割 rleSplitAt

仕様

{-
(***) :: (a -> c) -> (b -> d) -> (a, b) -> (c, d)
splitAt n xs == (fromRLE *** fromRLE) (rleSplitAt n (toRLE xs))
-}

実装

rleSplitAt :: Int -> RLE a -> (RLE a, RLE a)
rleSplitAt = \ case
    n+1 -> \ case
        h@(x,m+1):rs
            | n < m     -> ([(x,n+1)],(x,m-n):rs)
            | otherwise -> case rleSplitAt (n-m) rs of
                (as,bs)     -> (h:as,bs)
        _               -> ([],[])
    _   -> ([],)

テスト

{- ^
>>> let xs = "aaabcc"
>>> let expected n = splitAt n xs == (fromRLE *** fromRLE $ rleSplitAt n (toRLE xs))
>>> all expected [0 .. length xs]
True
-}

応用例 AtCoder Beginner Contest 413 C Large Queue

https://atcoder.jp/contests/abc413/tasks/abc413_c

type Query = [Int]
type Output = Int

solve :: [Query] -> [Output]
solve = snd . uncurry (mapAccumL psi) . foldr phi ([],[])
    where
        phi = \ case
            [1,c,x] -> first ((x, c) :) -- first  :: (a -> c) -> (a, b) -> (c, b)
            [2,k]   -> second (k :)     -- second :: (b -> c) -> (a, b) -> (a, c)
            _       -> error "invalid query"
        psi rle k = case rleSplitAt k rle of
            (rs,ss) -> (ss, foldl' (\ a (b,c) -> a + b * c) 0 rs)

追記（2025-07-10）

phi :: a -> b -> (a, c)が既知なら

snd (mapAccumL phi e xs) = unfoldr psi (e, xs)

となるpsiを以下のように構成できる。

psi :: (a, [b]) -> Maybe ((a, [b]), c)
psi (acc, xxs) = \ case
    x:xs -> case phi acc x of
        (acc', y) -> Just ((acc', xs), y)
    _    -> Nothing

ということは、solve は anamorphism と catamorphism をこの順で合成したもの、すなわち、metamorphism になっているということになりそう。recursion-schemesを使うと

import Data.Functor.Foldable

type Query = [Int]
type Output = Int

solve :: [Query] -> [Output]
solve = ana ψ . cata φ where
    φ = \ case
        Nil       -> ([],[])
        Cons a b -> case a of
            [1,c,x]   -> first ((x,c) :) b
            [2,k]     -> second (k :) b
            _         -> error "invalid query"
    ψ = \ case
        (_,[])   -> Nil
        (a,k:ks) -> case rleSplitAt k a of
            (b,c)    -> Cons (foldl (\ x (y,z) -> x + y * z) 0 b) (c,ks)

あ。AtCoderのHaskellのジャッジ環境では、recursion-schemesパッケージが含まれていないので、↑ を含むコードは CE になります。