Haskell基于数据构造器的列表简化函数优化实现问询

Great question! Your original approach (using list comprehensions or separate passes to sum A values and collect B elements) does indeed traverse the list twice. Let's fix that with a single traversal using a fold that tracks both the accumulated sum of A values and the collected B elements in one go.

First, let's recap what your two-pass implementation might look like for context:

data D = A Int | B Int deriving Show

-- Two-pass version (traverses the list twice)
simplifyTwoPass :: [D] -> [D]
simplifyTwoPass xs = 
  let totalA = sum [n | A n <- xs]
      bs = [b | b@(B _) <- xs]
  in if totalA == 0 then bs else A totalA : bs

Single-Pass Optimization with foldr #

The key is to use a fold that maintains a state tuple: (accumulatedASum, collectedBs). As we iterate through each element, we update this state accordingly—either adding to the A sum or appending a B element to our collection. Using foldr avoids needing to reverse the collected B list (unlike foldl), making it cleaner and more efficient for this use case.

Here's the optimized implementation:

simplify :: [D] -> [D]
simplify xs = 
  let (totalA, bs) = foldr updateState (0, []) xs
  in (if totalA /= 0 then [A totalA] else []) ++ bs
  where
    -- Update the state for each element in the list
    updateState :: D -> (Int, [D]) -> (Int, [D])
    updateState (A n) (acc, bs) = (acc + n, bs)
    updateState b@(B _) (acc, bs) = (acc, b : bs)

How It Works #

Let's walk through your example input [A 1, A 2, A 3, B 1, A 4, B 2]:

1. We start with an initial state (0, []).
2. Traversing from right to left (thanks to foldr):
   • Process B 2: state becomes (0, [B 2])
   • Process A 4: state becomes (4, [B 2])
   • Process B 1: state becomes (4, [B 1, B 2])
   • Process A 3: state becomes (7, [B 1, B 2])
   • Process A 2: state becomes (9, [B 1, B 2])
   • Process A 1: state becomes (10, [B 1, B 2])
3. Finally, we check if totalA (10) is non-zero, prepend A 10 to the collected B elements, and get the desired output: [A 10, B 1, B 2].

Edge Cases Handled #

• If there are no A elements: returns only the collected B elements.
• If there are no B elements: returns a single A with the total sum (or an empty list if the input was empty).
• Empty input: returns an empty list.

Why This Is Better #

Instead of traversing the list twice (once to sum As, once to collect Bs), we do everything in one O(n) pass. While the asymptotic time complexity is technically the same, this reduces practical overhead—especially for large lists—and is more idiomatic Haskell for stateful traversals of this kind.

内容的提问来源于stack exchange，提问作者Allan W