diff --git a/clustering.cabal b/clustering.cabal
--- a/clustering.cabal
+++ b/clustering.cabal
@@ -2,7 +2,7 @@
 -- documentation, see http://haskell.org/cabal/users-guide/
 
 name:                clustering
-version:             0.1.1
+version:             0.1.2
 synopsis:            High performance clustering algorithms
 description:
   Following clutering methods are included in this library:
diff --git a/src/AI/Clustering/Hierarchical.hs b/src/AI/Clustering/Hierarchical.hs
--- a/src/AI/Clustering/Hierarchical.hs
+++ b/src/AI/Clustering/Hierarchical.hs
@@ -41,7 +41,7 @@
 module AI.Clustering.Hierarchical
     ( Dendrogram(..)
     , size
-    , Metric(..)
+    , Linkage(..)
     , hclust
     , cutAt
     , flatten
@@ -63,16 +63,16 @@
 import AI.Clustering.Hierarchical.Types
 
 -- | Different hierarchical clustering schemes.
-data Metric = Single    -- ^ O(n^2) Single linkage, $d(A,B) = min_{a \in A, b \in B} d(a,b)$.
-            | Complete  -- ^ O(n^2) Complete linkage, $d(A,B) = max_{a \in A, b \in B} d(a,b)$.
-            | Average   -- ^ O(n^2) Average linkage or UPGMA, $d(A,B) = \frac{\sum_{a \in A}\sum_{b \in B}d(a,b)}{|A||B|}$.
-            | Weighted  -- ^ O(n^2) Weighted linkage.
-            | Ward      -- ^ O(n^2) Ward's method.
-            | Centroid  -- ^ O(n^3) Centroid linkage, not implemented.
-            | Median    -- ^ O(n^3) Median linkage, not implemented.
+data Linkage = Single    -- ^ O(n^2) Single linkage, $d(A,B) = min_{a \in A, b \in B} d(a,b)$.
+             | Complete  -- ^ O(n^2) Complete linkage, $d(A,B) = max_{a \in A, b \in B} d(a,b)$.
+             | Average   -- ^ O(n^2) Average linkage or UPGMA, $d(A,B) = \frac{\sum_{a \in A}\sum_{b \in B}d(a,b)}{|A||B|}$.
+             | Weighted  -- ^ O(n^2) Weighted linkage.
+             | Ward      -- ^ O(n^2) Ward's method.
+             | Centroid  -- ^ O(n^3) Centroid linkage, not implemented.
+             | Median    -- ^ O(n^3) Median linkage, not implemented.
 
 -- | Perform hierarchical clustering.
-hclust :: G.Vector v a => Metric -> v a -> DistFn a -> Dendrogram a
+hclust :: G.Vector v a => Linkage -> v a -> DistFn a -> Dendrogram a
 hclust method xs f = label <$> nnChain dists fn
   where
     dists = computeDists f xs
diff --git a/src/AI/Clustering/Hierarchical/Internal.hs b/src/AI/Clustering/Hierarchical/Internal.hs
--- a/src/AI/Clustering/Hierarchical/Internal.hs
+++ b/src/AI/Clustering/Hierarchical/Internal.hs
@@ -27,8 +27,12 @@
         | otherwise = go ds activeNodes $ c : chain
       where
         (c,d) = nearestNeighbor ds b a activeNodes
+
+        -- We always remove the node with smaller index. The other one will be
+        -- used to represent the merged result
         activeNodes' = M.insert hi (Branch (size1+size2) d c1 c2)
                      . M.delete lo $ activeNodes
+
         ds' = fn lo hi activeNodes ds
         c1 = M.findWithDefault undefined lo activeNodes
         c2 = M.findWithDefault undefined hi activeNodes
@@ -39,10 +43,16 @@
       where
         a = fst $ M.elemAt 0 activeNodes
         b = fst $ nearestNeighbor ds a (-1) activeNodes
+
     initSet = M.fromList . map (\i -> (i, Leaf i)) $ [0..n-1]
 {-# INLINE nnChain #-}
 
-nearestNeighbor :: DistanceMat -> Int -> Int -> M.Map Int (Dendrogram Int) -> (Int, Double)
+nearestNeighbor :: DistanceMat                  -- ^ distance matrix
+                -> Int                          -- ^ query
+                -> Int                          -- ^ this would be selected if
+                                                -- it achieves the minimal distance
+                -> M.Map Int (Dendrogram Int)
+                -> (Int, Double)
 nearestNeighbor dist i preference = M.foldlWithKey' f (-1,1/0)
   where
     f (x,d) j _ | i == j = (x,d)  -- skip
@@ -114,7 +124,7 @@
         d_lo_i <- UM.unsafeRead v $ idx n i lo
         d_hi_i <- UM.unsafeRead v $ idx n i hi
         UM.unsafeWrite v (idx n i hi) $
-            sqrt $ ((s1+s3)*d_lo_i + (s2+s3)*d_hi_i - s3*d_lo_hi) / (s1+s2+s3)
+            ((s1+s3)*d_lo_i + (s2+s3)*d_hi_i - s3*d_lo_hi) / (s1+s2+s3)
     return v
   where
     s1 = fromIntegral . size . M.findWithDefault undefined lo $ nodeset
diff --git a/src/AI/Clustering/KMeans.hs b/src/AI/Clustering/KMeans.hs
--- a/src/AI/Clustering/KMeans.hs
+++ b/src/AI/Clustering/KMeans.hs
@@ -1,6 +1,6 @@
 --------------------------------------------------------------------------------
 -- |
--- Module      :  $Header$
+-- Module      :  AI.Clustering.KMeans
 -- Copyright   :  (c) 2015 Kai Zhang
 -- License     :  MIT
 -- Maintainer  :  kai@kzhang.org
@@ -9,45 +9,59 @@
 --
 -- Kmeans clustering
 --------------------------------------------------------------------------------
-{-# LANGUAGE FlexibleContexts  #-}
-
 module AI.Clustering.KMeans
-    ( kmeans
+    ( KMeans(..)
+    , kmeans
     , kmeansWith
-    , forgyMethod
+
+    -- * Initialization methods
+    , Initialization(..)
+
+    , decode
     ) where
 
 import Control.Monad (forM_)
 import Control.Monad.Primitive (PrimMonad, PrimState)
-import qualified Data.Matrix.Generic as M
-import qualified Data.Matrix.Generic.Mutable as MM
+import qualified Data.Matrix.Unboxed as MU
+import qualified Data.Matrix.Unboxed.Mutable as MM
 import Data.Ord (comparing)
+import qualified Data.Vector as V
+import qualified Data.Vector.Mutable as VM
 import qualified Data.Vector.Unboxed as U
 import qualified Data.Vector.Unboxed.Mutable as UM
-import qualified Data.Vector.Generic as G
-import qualified Data.Vector.Generic.Mutable as GM
 import Data.List (minimumBy, nub)
 import System.Random.MWC (uniformR, Gen)
 
+-- | Results from running kmeans
+data KMeans = KMeans
+    { _clusters :: U.Vector Int     -- ^ A vector of integers (0 ~ k-1)
+                                    -- indicating the cluster to which each
+                                    -- point is allocated.
+    , _centers :: MU.Matrix Double  -- ^ A matrix of cluster centres.
+    } deriving (Show)
+
 -- | Lloyd's algorithm, also known as K-means algorithm
-kmeans :: (G.Vector v Double, G.Vector v Int, Eq (v Int), Eq (v Double), PrimMonad m)
+kmeans :: PrimMonad m
        => Gen (PrimState m)
+       -> Initialization
        -> Int                           -- ^ number of clusters
-       -> M.Matrix v Double             -- ^ each row represents a point
-       -> m (v Int, M.Matrix v Double)  -- ^ membership vector
-kmeans g k mat = do
-    initial <- forgyMethod g k mat
-    return $ kmeansWith mat initial
+       -> MU.Matrix Double             -- ^ each row represents a point
+       -> m KMeans
+kmeans g method k mat = do
+    initial <- case method of
+        Forgy -> forgy g k mat
+        _ -> undefined
+    return $ kmeansWith initial mat
 {-# INLINE kmeans #-}
 
 -- | Lloyd's algorithm, also known as K-means algorithm
-kmeansWith :: (G.Vector v Double, G.Vector v Int, Eq (v Int))
-           => M.Matrix v Double           -- ^ initial set of k centroids
-           -> M.Matrix v Double           -- ^ each row represents a point
-           -> (v Int, M.Matrix v Double)  -- ^ membership vector and centroids
-kmeansWith initial mat | d /= M.cols initial || k > n = error "check input"
-                       | otherwise = loop initial G.empty
+kmeansWith :: MU.Matrix Double   -- ^ initial set of k centroids
+           -> MU.Matrix Double   -- ^ each row represents a point
+           -> KMeans
+kmeansWith initial mat | d /= MU.cols initial || k > n = error "check input"
+                       | otherwise = KMeans member centers
   where
+    (member, centers) = loop initial U.empty
     loop means membership
         | membership' == membership = (membership, means)
         | otherwise = loop (update membership') membership'
@@ -55,63 +69,78 @@
         membership' = assign means
 
     -- Assignment step
-    assign means = G.generate n $ \i ->
-        let x = M.takeRow mat i
-        in fst $ minimumBy (comparing snd) $ zip [0..k-1] $ map (dist x) $ M.toRows means
+    assign means = U.generate n $ \i ->
+        let x = MU.takeRow mat i
+        in fst $ minimumBy (comparing snd) $ zip [0..k-1] $ map (dist x) $ MU.toRows means
 
-    --  Update step
+    -- Update step
     update membership = MM.create $ do
         m <- MM.replicate (k,d) 0.0
         count <- UM.replicate k (0 :: Int)
         forM_ [0..n-1] $ \i -> do
-            let x = membership G.! i
-            GM.unsafeRead count x >>= GM.unsafeWrite count x . (+1)
+            let x = membership U.! i
+            UM.unsafeRead count x >>= UM.unsafeWrite count x . (+1)
             forM_ [0..d-1] $ \j ->
-                MM.unsafeRead m (x,j) >>= MM.unsafeWrite m (x,j) . (+ mat M.! (i,j))
+                MM.unsafeRead m (x,j) >>= MM.unsafeWrite m (x,j) . (+ mat MU.! (i,j))
         -- normalize
         forM_ [0..k-1] $ \i -> do
-            c <- GM.unsafeRead count i
+            c <- UM.unsafeRead count i
             forM_ [0..d-1] $ \j ->
                 MM.unsafeRead m (i,j) >>= MM.unsafeWrite m (i,j) . (/fromIntegral c)
         return m
 
-    dist :: G.Vector v Double => v Double -> v Double -> Double
-    dist xs = G.sum . G.zipWith (\x y -> (x - y)**2) xs
+    dist xs = U.sum . U.zipWith (\x y -> (x - y)**2) xs
 
-    n = M.rows mat
-    k = M.rows initial
-    d = M.cols mat
+    n = MU.rows mat
+    k = MU.rows initial
+    d = MU.cols mat
 {-# INLINE kmeansWith #-}
 
--- * Initialization methods
+-- | Different initialization methods
+data Initialization = Forgy    -- ^ The Forgy method randomly chooses k unique
+                               -- observations from the data set and uses these
+                               -- as the initial means
+                    | KMeansPP -- ^ K-means++ algorithm, not implemented.
 
--- | The Forgy method randomly chooses k unique observations from the data set and uses
--- these as the initial means
-forgyMethod :: (PrimMonad m, G.Vector v a, Eq (v a))
-            => Gen (PrimState m)
-            -> Int                 -- number of clusters
-            -> M.Matrix v a        -- data
-            -> m (M.Matrix v a)
-forgyMethod g k mat | k > n = error "k is larger than sample size"
-                    | otherwise = iter
+forgy :: PrimMonad m
+      => Gen (PrimState m)
+      -> Int                 -- number of clusters
+      -> MU.Matrix Double    -- data
+      -> m (MU.Matrix Double)
+forgy g k mat | k > n = error "k is larger than sample size"
+              | otherwise = iter
   where
     iter = do
         vec <- sample g k . U.enumFromN 0 $ n
-        let xs = map (M.takeRow mat) . G.toList $ vec
+        let xs = map (MU.takeRow mat) . U.toList $ vec
         if length (nub xs) == length xs
-           then return . M.fromRows $ xs
+           then return . MU.fromRows $ xs
            else iter
-    n = M.rows mat
-{-# INLINE forgyMethod #-}
+    n = MU.rows mat
+{-# INLINE forgy #-}
 
 -- random select k samples from a population
-sample :: (PrimMonad m, G.Vector v a) => Gen (PrimState m) -> Int -> v a -> m (v a)
+sample :: PrimMonad m => Gen (PrimState m) -> Int -> U.Vector Int -> m (U.Vector Int)
 sample g k xs = do
-    v <- G.thaw xs
+    v <- U.thaw xs
     forM_ [0..k-1] $ \i -> do
         j <- uniformR (i, lst) g
-        GM.unsafeSwap v i j
-    G.unsafeFreeze . GM.take k $ v
+        UM.unsafeSwap v i j
+    U.unsafeFreeze . UM.take k $ v
   where
-    lst = G.length xs - 1
+    lst = U.length xs - 1
 {-# INLINE sample #-}
+
+-- | Assign data to clusters based on KMeans result
+decode :: KMeans -> [a] -> [[a]]
+decode result xs = V.toList $ V.create $ do
+    v <- VM.replicate n [] 
+    forM_ (zip (U.toList membership) xs) $ \(i,x) ->
+        VM.unsafeRead v i >>= VM.unsafeWrite v i . (x:)
+    return v
+  where
+    membership = _clusters result
+    n = U.maximum membership + 1
+
+-- | Compute within-cluster sum of squares
+--withinSS :: Matrix