如何优化这个Haskell函数？

Question

我需要在调色板ps中找到与给定颜色p最接近的颜色。如何使函数nearestColor尽可能快，而不改变Pixel8或PixelRGB8的类型。到目前为止，我已经尝试了内联。

import qualified Data.Vector as V

type Pixel8 = Word8    

data PixelRGB8 = PixelRGB8 {-# UNPACK #-} !Pixel8 -- Red
                           {-# UNPACK #-} !Pixel8 -- Green
                           {-# UNPACK #-} !Pixel8 -- Blue
           deriving (Eq, Ord, Show)

nearestColor :: PixelRGB8 -> Vector PixelRGB8 -> PixelRGB8
nearestColor p ps = snd $ V.minimumBy comp ds
  where
    ds = V.map (\px -> (dist2Px px p, px)) ps
    comp a b = fst a `compare` fst b

dist2Px :: PixelRGB8 -> PixelRGB8 -> Int
dist2Px (PixelRGB8 r1 g1 b1) (PixelRGB8 r2 g2 b2) = dr*dr + dg*dg + db*db
  where
    (dr, dg, db) =
      ( fromIntegral r1 - fromIntegral r2
      , fromIntegral g1 - fromIntegral g2
      , fromIntegral b1 - fromIntegral b2 )

- Yofe

---

我们能看到dist2Px的代码吗？ - daniel gratzer

你目前尝试了哪些方法？你可以尝试对其进行分析，可以将本地函数“ds”和“comp”进行内联，可以编译核心并搜索瓶颈，然后编译为汇编语言并确保您有紧密的循环。 - bheklilr

我怀疑只使用距离构建向量，并使用 minIndex 在原始向量中查找结果可能会更快：ps ! V.minIndex (V.map (\px -> dist2Px px p) ps)。 - Tarmil

@Tamil - 这不会加快速度，但更加简洁，谢谢。 - Yofe

你到底想要优化什么：针对单个函数调用的时间，还是针对不同颜色但相同的调色板进行大量调用的平均时间？ - leftaroundabout

我想要针对不同颜色但相同调色板的大量调用进行平均优化。 - Yofe

Answer 1

如果您想使用单个调色板并请求不同的颜色，我建议先翻转您的签名：

type Palette = V.Vector PixelRGB8
nearestColor :: Palette -> PixelRGB8 -> PixelRGB8

这有助于部分应用，并允许记忆调色板配置。

接下来，您需要将调色板重新存储在适合快速查找的数据结构中。由于您基本上对 ℝ³ 中的欧几里得距离感兴趣（顺便说一句，这并不是用于颜色比较的理想选择），这是一个非常普遍的问题。经典的结构是 k-d 树，长期以来一直被用于最近邻搜索。有足够方便的 Haskell 库 可供使用：

import qualified Data.Trees.KdTree a s KD

instance KD.Point PixelRGB where
  dimension _ = 3
  coord 0 (PixelRGB r _ _) = fromIntegral r
  coord 1 (PixelRGB _ g _) = fromIntegral g
  coord 2 (PixelRGB _ _ b) = fromIntegral b
  dist2 = fromIntegral . dist2Px

然后我们可以将调色板转换成这样的树形结构：

type FastPalette = KD.KdTree PixelRGB8
accelPalette :: Palette -> FastPalette
accelPalette = KD.fromList . V.toList

最后只需使用库提供的下一个邻居搜索：

nearestColor palette = fromJust . KD.nearestNeighbor fpal
 where fpal = accelPalette palette

- leftaroundabout

---

我点赞了，因为这是一个很棒的答案。不幸的是，我只能使用 Haskell 平台中的软件包。顺便说一下，还有一些类型是 coord1 ... = fromIntegral g 等。 - Yofe

当然，你也可以重新实现_k_-d树，这并不难。 - leftaroundabout