P23 - 从列表随机选 N 个元素
Extract a given number of randomly selected elements from a list
官方模块:
Problems.P23核心函数:randomSelect
← P22 生成整数范围 | P24 从 1..M 中随机选 N 个不同数 →
题目描述
从列表中随机选择 个不重复的元素。需要接受一个随机数生成器并返回更新后的生成器——这是 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
先把前 个元素放入蓄水池。读到第 个后续元素时,在 到 之间随机选择位置;只有位置落在蓄水池内才执行替换。这样只需扫描一次列表,并且每个元素最终被选中的概率相同。
方法三:基于洗牌的固定抽样
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