從快速排序 Quicksort 到快速選擇 Quickselect

函數式編程(functional programming)相對於命令式編程(imperative programming),特別適合使用程式推導(program derivation):有別於傳統演算法設計依靠直覺經驗跟觀察,系統化地由一個描述性的演算法作為問題的定義(specification),再由各種演算法構件間的關係推導出一個等價但時間跟空間複雜度都較低的演算法。

一個最有名例子是「最大區段和(maximum segment sum)」 的問題,這問題經典隨處可見,今天來寫一個更簡單的問題:在一個未排序的串列(list)中,如何選出第 k 小的元素呢?

我們把問題表達成以下的 Haskell 程式作為問題的定義:

select :: (Ord a) => Int -> [a] -> a
select k = (!! k) . sort

意即:排序後第 k 個的元素即是原本串列中第 k 小的元素;如果 k 大於給定串列的長度,則沒有定義。時間複雜度上,對串列作 index 需要 O(n) 的時間,而以比較為主的排序演算法,平均時間複雜度至少是 O(nlog n),因此現有的 select 需要 O(nlog n) 。能不能做得更好呢?

假設用 Quicksort 作排序,在 Haskell 寫下一個非 in-place 的版本很容易:

sort :: (Ord a) => [a] -> [a]
sort []     = []
sort (x:xs) = sort ys ++ [x] ++ sort zs 
  where
    (ys, zs) = (filter (<x) xs, filter (not . (<x)) xs)

而 index 操作在 Haskell 則是定義為:

(!!) :: [a] -> Int -> a
(x:xs) !! 0 = x
(x:xs) !! k = xs !! (k-1)
_      !! _ = undefined

從定義可以得知若 k 為負或是大於串列的長度時,程式沒有定義。現在考慮串列的兩種情況:空串列跟一個以上元素的串列。首先是空串列的情況:

  sort [] !! k
= { 根據 sort 的定義 } 
  [] !! k
= { 根據 [] !! k 的定義 }
  無定義

非空串列的情況:

  sort (x:xs) !! k
= { 根據 sort 的定義 }
  (sort ys ++ [x] ++ sort zs) !! k
    where
      (ys, zs) = (filter (<x) xs, filter (not . (<x)) xs)

其中把 sort 對非空串列的定義展開後,得到含有串接 (++) 的表達式。假設我們知道前段串列的長度,那麼我們可以比較 index 的位置跟長度,來決定要取哪段串列。在 Haskell 有個 compare 函數直接說明兩個數字間式小於(LT),等於(EQ)還是大於(EQ)我們用這個函數比較大小,前述 (++)(!!) 之間的等式可以表達成:

(xs ++ [y] ++ ys) !! k =
  case compare k n of 
    LT -> xs !! k
    EQ -> y
    GT -> ys !! (k - n - 1)
  where n = length xs

(證明留給讀者玩玩看)有了以上的等式之後,我們可以繼續把卡住的部分展開:

  (sort ys ++ [x] ++ sort zs) !! k
    where
      (ys, zs) = (filter (<x) xs, filter (not . (<x)) xs)

= case compare k n of
    LT -> (sort ys) !! k
    EQ -> x 
    GT -> (sort zs) !! (n - k -1)
      where
        (ys, zs) = (filter (<x) xs, filter (not . (<x)) xs)
        n = length ys

= { 根據 select 的定義 }
  case compare k n of 
    LT -> select k ys 
    EQ -> x
    GT -> select (k - n - 1) zs 
      where
        (ys, zs) = (filter (<x) xs, filter (not . (<x)) xs)
        n = length ys

綜合前面兩種情況,就得到了以下等式:

select k (x:xs) =
  case compare k n of
    LT -> select k ys 
    EQ -> x
    GT -> select (k - n - 1) zs 
      where
        (ys, zs) = (filter (<x) xs, filter (not . (<x)) xs)
        n = length ys

幾乎是我們熟知的 Quickselect了。但我們一共在串列 xs 上套用 filter 函數兩次,另外再計算 ys 的長度一次,總共走過串列三次,能不能只走一次呢?

在 Haskell 有個 partition 函數其計算等價於

partition p xs = (filter (<x) xs, filter (not . (<x)) xs)

只需要走過串列一次。有沒有辦法再合併計算 filter (<x) xs 的長度呢?我們可以透過 Tupling 的技巧來做到:

對任意函數 f, g 以及元素 a, b 以及串列 xs,以下等式成立:

h x (y, z) (f x y, g x z) 則

(foldr f a xs, foldr g b xs) = foldr h (a, b) xs

首先將 (length (filter p xs), filter p xs, filter (not . p) xs)foldr 表達(留給讀者),再作 tupling 兩次得到同時計算 partition 跟前段符合 p 的串列長度:

partition' p = foldr op [] 
  where op x (n, ys, zs) | p x       = (n+1, x:ys, zs)
                         | otherwise = (n, ys, x:xs)

於是整理目前已知,就得到最終的 Quickselect:

select k (x:xs) =
  case compare k n of
    LT -> select k ys 
    EQ -> x
    GT -> select (k - n - 1) zs 
      where
        (n, ys, zs) = partition' (<x) xs 

partition' p = foldr op [] 
  where op x (n, ys, zs) | p x       = (n+1, x:ys, zs)
                         | otherwise = (n, ys, x:xs)

最後用的 Haskell 程式碼含測試

{-# LANGUAGE TypeApplications #-}
import Test.QuickCheck
import Data.List (sort)
selectOrigin k = (!! k) . sort
select :: (Ord a) => Int -> [a] -> a
select k (x:xs) = case compare k n of
LT -> select k ys
EQ -> x
GT -> select (k-n-1) zs
where
(n, ys, zs) = partition' (< x) xs
partition' :: (a -> Bool) -> [a] -> (Int, [a], [a])
partition' p = foldr op (0, [], [])
where
op x (n, ys, zs)
| p x = (1+n, x:ys, zs)
| otherwise = (n, ys, x:zs)
prop1 :: Ord a => NonEmptyList a -> Int -> Bool
prop1 (NonEmpty xs) k =
let k' = k `mod` length xs
in selectOrigin k' xs == select k' xs
prop2 :: [Int] -> Int -> [Int] -> Int -> Bool
prop2 xs y ys n =
let m = length xs
n' = n `mod` (m + length (y:ys))
in (xs ++ y:ys) !! n' ==
case compare n' m of
LT -> xs !! n'
EQ -> y
GT -> ys !! (n' - m - 1)
main :: IO ()
main = quickCheck (prop1 @Int)
>> quickCheck prop2
view raw Quickselect.hs hosted with ❤ by GitHub

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s