BKTree implementation

imageduplicates.cabal

@@ -18,12 +18,13 @@ cabal-version:       >=1.10
 executable imageduplicates
   main-is:             Main.hs
   ghc-options:         -Wall -threaded -rtsopts -O2
-  -- other-modules:
+  other-modules:       Data.BKTree
   -- other-extensions:
   build-depends:       base >=4.11 && <4.12
                      , JuicyPixels
                      , JuicyPixels-extra
                      , bytestring
+                     , either
                      , generic-lens
                      , lens
                      , mtl
@@ -31,6 +32,7 @@ executable imageduplicates
                      , pipes
                      , pipes-files
                      , pipes-safe
+                     , recursion-schemes
                      , text
                      , transformers
   hs-source-dirs:      src

src/Data/BKTree.hs

@@ -0,0 +1,54 @@
+{-# LANGUAGE DeriveFoldable    #-}
+{-# LANGUAGE DeriveFunctor     #-}
+{-# LANGUAGE DeriveGeneric     #-}
+{-# LANGUAGE DeriveTraversable #-}
+{-# LANGUAGE KindSignatures    #-}
+{-# LANGUAGE LambdaCase        #-}
+{-# LANGUAGE TemplateHaskell   #-}
+{-# LANGUAGE TypeFamilies      #-}
+{-# LANGUAGE ScopedTypeVariables      #-}
+module Data.BKTree where
+
+import           Data.Functor.Foldable.TH
+import           Data.Functor.Foldable
+import           GHC.Generics             (Generic)
+
+-- Point for testing purposes
+data Point = Point Int Int deriving Show
+
+instance Metric Point where
+  distance (Point p1 p2) (Point q1 q2) = abs (p1 - q1) + abs (p2 - q2)
+
+class Metric a where
+  distance :: a -> a -> Int
+
+data BKTree a = Empty
+              | Node a [(Int, BKTree a)] deriving (Show, Generic)
+
+makeBaseFunctor ''BKTree
+
+insert :: Metric a => a -> BKTree a -> BKTree a
+insert a = \case
+  Empty -> Node a []
+  Node b children ->
+    let newDistance = distance a b
+    in Node b (addChild newDistance children)
+  where
+    addChild d = \case
+      [] -> (d, insert a Empty) : []
+      (d',child):children | d == d' -> (d', insert a child) : children
+                          | otherwise -> (d',child) : addChild d children
+
+
+search :: forall a. Metric a => Int -> a -> BKTree a -> [a]
+search n a tree = cata alg tree
+  where
+    alg :: BKTreeF a [a] -> [a]
+    alg = \case
+      EmptyF -> []
+      NodeF x children ->
+        let thisDistance = distance a x
+            upper = thisDistance + n
+            lower = thisDistance - n
+            filteredChildren = concat [xs | (d, xs) <- children, d <= upper, d >= lower]
+        in if thisDistance <= n then x : filteredChildren else filteredChildren