用State Monad重构不纯递归：确认Haskell代码对应Python实现

问题背景 #

我正在分析Advent of Code 2022第14题的Python单行不纯递归解决方案：

def s(x,y):
    if y > h: return True
    if (x, y) in m: return False
    return next((r for d in (0,-1,1) if (r:=s(x+d,y+1))), None) or m.add((x,y))

已知可通过显式传递集合m将其改为纯函数，但希望借助State Monad避免额外参数传递。找到一段疑似对应的Haskell实现：

fill :: (MArray a Bool (ST s), Ix i, Num i, Show i) => a (i, i) Bool -> i -> ST s (Int, Int)
fill blocks maxY = do
    counterAtMaxY <- newSTRef Nothing
    counter <- newSTRef 0
    let fill' (x, y) = readArray blocks (x, y) >>= flip bool (pure ()) do
            when (y == maxY) $ readSTRef counterAtMaxY >>= maybe
                (readSTRef counter >>= writeSTRef counterAtMaxY . Just) (const $ pure ())
            when (y <= maxY) $ fill' (x, y + 1) >> fill' (x - 1, y + 1) >> fill' (x + 1, y + 1)
            writeArray blocks (x, y) True >> modifySTRef' counter (+ 1)
    fill' (500, 0)
    counterAtMaxY <- readSTRef counterAtMaxY
    counter <- readSTRef counter
    pure (fromMaybe counter counterAtMaxY, counter)

需要确认：

1. 该Haskell代码是否为上述Python代码的移植
2. 详细讲解递归逻辑
3. 将代码中同时处理问题第二部分的逻辑分离，避免混淆坐标与结果元组

回答 #

1. 是否为Python代码的移植？ #

是，这段Haskell代码本质是Python实现的移植，同时扩展了功能（一次性处理题目两部分）。两者核心逻辑完全对齐：

• 模拟沙粒从(500,0)下落，优先尝试向下→左下→右下的路径
• 用可变状态记录已占据的位置（Python用集合m，Haskell用ST Monad下的可变数组blocks）
• 只有当三个方向都被阻挡时，沙粒才会停留在当前位置

差异点在于Haskell代码额外加入了计数器，直接统计题目两部分的结果，而Python代码仅通过副作用修改集合，返回值仅用于控制递归流程。

2. 递归逻辑详解 #

Python原逻辑拆解

• 终止条件1：y > h（沙粒超出底部边界，直接流失），返回True表示该路径无法停留沙粒
• 终止条件2：(x,y)在m中（位置已被占据），返回False表示无法进入该位置
• 递归尝试：依次向下、左下、右下递归调用，只要其中一个方向返回True（沙粒从该方向流失），当前沙粒也会流失；若三个方向都返回False（全部被阻挡），则将(x,y)加入m（沙粒停留），最终返回None（Python中set.add返回None）

Haskell递归逻辑拆解

Haskell的fill'函数借助ST Monad维护可变状态，逻辑对应如下：

1. 读取当前坐标(x,y)的状态：如果已被占据（blocks值为True），则直接跳过
2. 若当前坐标未被占据，进入核心逻辑：
   • 检查y == maxY（题目第一部分的底部边界）：如果是且未记录过该状态，保存当前计数器值（即第一部分结果：沙粒开始流失时的总数）
   • 若y <= maxY（未超出边界），按向下→左下→右下顺序递归处理三个方向（确保沙粒优先走最自然的下落路径）
   • 标记当前坐标为已占据，并将计数器+1（统计停留的沙粒数，对应第二部分结果：直到顶部被堵住的总数）
3. 递归起点为(500,0)，即沙粒的初始下落位置

3. 分离两部分逻辑的改进方案 #

原代码将两部分结果统计混在一起，状态共享容易混淆。以下是拆分后的实现，将两部分逻辑完全独立：

import Control.Monad.ST
import Data.Array.MArray
import Data.Ix
import Data.STRef

-- 计算题目第一部分结果：沙粒开始从底部流失时的总数
fillPart1 :: (MArray a Bool (ST s), Ix i, Num i) => a (i, i) Bool -> i -> ST s Int
fillPart1 blocks maxY = do
    counter <- newSTRef 0
    let fill' (x, y) = readArray blocks (x, y) >>= flip bool (pure ()) do
            if y > maxY
                then pure ()  -- 沙粒流失，不计数
                else do
                    fill' (x, y + 1)
                    fill' (x - 1, y + 1)
                    fill' (x + 1, y + 1)
                    writeArray blocks (x, y) True
                    modifySTRef' counter (+1)
    fill' (500, 0)
    readSTRef counter

-- 计算题目第二部分结果：直到顶部(500,0)被堵住时的总数
fillPart2 :: (MArray a Bool (ST s), Ix i, Num i) => a (i, i) Bool -> i -> ST s Int
fillPart2 blocks maxY = do
    counter <- newSTRef 0
    let floorY = maxY + 1  -- 第二部分的虚拟底部
        fill' (x, y) = readArray blocks (x, y) >>= flip bool (pure ()) do
            if y == floorY
                then do
                    writeArray blocks (x, y) True
                    modifySTRef' counter (+1)
                else do
                    fill' (x, y + 1)
                    fill' (x - 1, y + 1)
                    fill' (x + 1, y + 1)
                    writeArray blocks (x, y) True
                    modifySTRef' counter (+1)
    fill' (500, 0)
    readSTRef counter

-- 统一入口：初始化数组并计算两部分结果
solveDay14 :: Ix i => [(i, i)] -> i -> (Int, Int)
solveDay14 rocks maxY = runST $ do
    -- 初始化数组边界，预留足够空间避免越界
    let bounds = ((500 - maxY - 1, 0), (500 + maxY + 1, maxY + 1))
    -- 计算第一部分
    blocks1 <- newArray bounds False
    mapM_ (\pos -> writeArray blocks1 pos True) rocks
    part1 <- fillPart1 blocks1 maxY
    -- 重新初始化数组计算第二部分（避免状态污染）
    blocks2 <- newArray bounds False
    mapM_ (\pos -> writeArray blocks2 pos True) rocks
    part2 <- fillPart2 blocks2 maxY
    pure (part1, part2)

改进说明

• 职责分离：fillPart1和fillPart2分别处理题目两部分，各自维护独立计数器，逻辑清晰无混淆
• 边界处理：fillPart1严格遵循原Python逻辑，沙粒超出边界时直接终止；fillPart2新增虚拟底部maxY+1，模拟题目第二部分的无限地面
• 状态隔离：通过重新初始化数组避免两部分计算互相干扰，符合纯函数的设计原则

内容的提问来源于stack exchange，提问作者Tom Huntington