P23 - 从列表随机选 N 个元素

2026-09-07 00:00    #Haskell   #99题   #列表  

P23 - 从列表随机选 N 个元素

Extract a given number of randomly selected elements from a list

官方模块:Problems.P23 核心函数:randomSelect


← P22 生成整数范围 | P24 从 1..M 中随机选 N 个不同数 →


题目描述

从列表中随机选择 NN 个不重复的元素。需要接受一个随机数生成器并返回更新后的生成器——这是 Haskell 纯函数随机数的标准做法。

函数签名

1import System.Random (RandomGen)
2
3randomSelect :: RandomGen g => [a] -> Int -> g -> ([a], g)

RandomGen g 是类型类约束,表示任何随机数生成器都可以。

实现

方法一:随机移除(选中的拿走)

 1import System.Random (RandomGen, randomR)
 2
 3randomSelect :: RandomGen g => [a] -> Int -> g -> ([a], g)
 4randomSelect xs count gen = go xs (max 0 count) gen []
 5  where
 6    go [] _ g picked = (reverse picked, g)
 7    go _  0 g picked = (reverse picked, g)
 8    go pool k g picked =
 9      let (i, g')       = randomR (0, length pool - 1) g
10          (x, pool')    = removeAt0 i pool
11      in go pool' (k - 1) g' (x : picked)
12
13removeAt0 :: Int -> [a] -> (a, [a])
14removeAt0 i xs =
15  case splitAt i xs of
16    (front, x : back) -> (x, front ++ back)
17    _ -> error "removeAt0: index out of bounds"

核心思路:每次从剩余池中随机选一个元素,将其从池中移除并加入结果列表。

注意:randomR (0, length pool - 1) 在空列表时会除零。好在 pool 为空时,前面的 go [] 模式已经截住了。

方法二:蓄水池抽样

 1import System.Random (RandomGen, randomR)
 2
 3randomSelect :: RandomGen g => [a] -> Int -> g -> ([a], g)
 4randomSelect xs count gen
 5  | count <= 0 = ([], gen)
 6  | otherwise  = go size reservoir rest gen
 7  where
 8    size = min count (length xs)
 9    (reservoir, rest) = splitAt size xs
10
11    go _ picked [] g = (picked, g)
12    go seen picked (x:remaining) g =
13      let (slot, g') = randomR (0, seen) g
14          picked'
15            | slot < size = replaceAt slot x picked
16            | otherwise   = picked
17      in go (seen + 1) picked' remaining g'
18
19replaceAt :: Int -> a -> [a] -> [a]
20replaceAt i x xs = front ++ x : drop 1 back
21  where
22    (front, back) = splitAt i xs

先把前 NN 个元素放入蓄水池。读到第 ii 个后续元素时,在 00ii 之间随机选择位置;只有位置落在蓄水池内才执行替换。这样只需扫描一次列表,并且每个元素最终被选中的概率相同。

方法三:基于洗牌的固定抽样

 1import System.Random (RandomGen, randomR)
 2
 3randomSelect :: RandomGen g => [a] -> Int -> g -> ([a], g)
 4randomSelect xs k gen = (take k shuffled, gen'')
 5  where
 6    (shuffled, gen'') = fisherYates xs gen
 7
 8fisherYates :: RandomGen g => [a] -> g -> ([a], g)
 9fisherYates [] gen = ([], gen)
10fisherYates xs gen = go (length xs) xs gen
11  where
12    go 0 acc gen = (acc, gen)
13    go n acc gen =
14      let (i, gen') = randomR (0, n - 1) gen
15          swapped   = swap i (n - 1) acc
16      in go (n - 1) swapped gen'
17
18swap :: Int -> Int -> [a] -> [a]
19swap i j xs
20  | i == j    = xs
21  | otherwise =
22      [ if k == i then xj else if k == j then xi else x
23      | (k, x) <- zip [0..] xs
24      ]
25  where
26    xi = xs !! i
27    xj = xs !! j

先对整个列表做 Fisher-Yates 洗牌,再取前 K 个。适合 K 接近列表长度的场景。

方法对比

方法特点效率
随机移除每轮从候选池删除一个元素O(N × len xs)
蓄水池抽样单次扫描,只保留 N 个元素O(len xs × N),数组实现可到 O(len xs)
先洗牌再取得到完整随机排列后截取当前列表实现 O(len xs²)

测试

 1>>> import System.Random (mkStdGen)
 2>>> let (picked, _) = randomSelect "abcdefgh" 3 (mkStdGen 111)
 3>>> length picked
 43
 5>>> let (picked, _) = randomSelect [1..5] 10 (mkStdGen 1)
 6>>> length picked    -- 最多只有 5 个可选
 75
 8>>> let (picked, _) = randomSelect [1..20] 5 (mkStdGen 7)
 9>>> all (`elem` [1..20]) picked && length picked == 5
10True

参考