packages feed

hanalyze-0.2.0.0: src/Hanalyze/Math/Hungarian.hs

{-# LANGUAGE BangPatterns #-}
-- |
-- Module      : Hanalyze.Math.Hungarian
-- Description : Hungarian (Kuhn-Munkres) 法による正方割当問題の最小コスト解 (O(n³))
-- Copyright   : (c) 2026 Aelysce Project (Toshiaki Honda)
-- License     : BSD-3-Clause
--
-- Hungarian (Kuhn-Munkres) アルゴリズムによる正方割当問題の最小コスト解。
--
-- ## 入出力
--
-- 入力: コスト行列 C (n × n、 各成分は実数、 inf 不可)。
-- 出力: 行 i に割当てる列 j からなる長さ n のベクトル @assignment[i] = j@。
-- 目的: Σᵢ C[i, assignment[i]] を最小化、 かつ assignment が **全単射**。
--
-- ## 実装
--
-- e-maxx の "Hungarian algorithm in O(V³)" 系統 (Jonker-Volgenant の
-- shortest augmenting path 方式)。 双対変数 u, v と potential を保持して
-- 1 行ずつ augment する。 ST + mutable Vector で内部状態を管理し、 純関数
-- 'hungarianMin' として API 公開する。
--
-- ## 用途
--
-- ICA-LiNGAM (Shimizu 2006) の行/列順列下三角化で、 W 行列の対角成分絶対値
-- を最大化する割当を求めるのに使う。 コスト C[i, j] = 1 / (|W[i, j]| + ε)
-- で 'hungarianMin' を呼ぶと、 グリーディと違って大域最適解が得られる。
-- p > 10 でグリーディが劣化するケースを救う。
--
-- ## 計算量
--
-- O(n³)。 n ≤ 200 程度では実用上問題なし (測定: n=100 で数十 ms オーダー、
-- 計測値ではなく目安)。
module Hanalyze.Math.Hungarian
  ( hungarianMin
  ) where

import           Control.Monad               (forM_, unless, when)
import           Control.Monad.ST            (ST, runST)
import           Data.STRef
import qualified Data.Vector.Unboxed         as VU
import qualified Data.Vector.Unboxed.Mutable as MV
import qualified Numeric.LinearAlgebra       as LA

-- ===========================================================================
-- 公開 API
-- ===========================================================================

-- | 正方コスト行列 C (n × n) に対する最小コスト割当。
--   戻り値 @v@ は @v VU.! i = j@ で「行 i が列 j に割当てられる」 意味。
hungarianMin :: LA.Matrix Double -> VU.Vector Int
hungarianMin cost
  | n == 0    = VU.empty
  | otherwise = runST (runHungarian n cost)
  where
    n = LA.rows cost

-- ===========================================================================
-- 内部実装 (ST monad、 1-indexed の慣例で size n+1 配列を確保)
-- ===========================================================================

runHungarian :: Int -> LA.Matrix Double -> ST s (VU.Vector Int)
runHungarian n cost = do
  let !inf = 1.0e300 :: Double
  u   <- MV.replicate (n + 1) (0 :: Double)
  v   <- MV.replicate (n + 1) (0 :: Double)
  p   <- MV.replicate (n + 1) (0 :: Int)     -- p[j] = 列 j に割当てた行
  way <- MV.replicate (n + 1) (0 :: Int)

  forM_ [1 .. n] $ \i -> do
    MV.write p 0 i
    j0Ref <- newSTRef (0 :: Int)
    minv  <- MV.replicate (n + 1) inf
    used  <- MV.replicate (n + 1) False

    let -- shortest-path-tree 拡張 1 ステップ
        step = do
          j0 <- readSTRef j0Ref
          MV.write used j0 True
          i0 <- MV.read p j0
          deltaRef <- newSTRef inf
          j1Ref    <- newSTRef (0 :: Int)
          forM_ [1 .. n] $ \j -> do
            isU <- MV.read used j
            unless isU $ do
              ui0 <- MV.read u i0
              vj  <- MV.read v j
              let !cur = LA.atIndex cost (i0 - 1, j - 1) - ui0 - vj
              mj <- MV.read minv j
              when (cur < mj) $ do
                MV.write minv j cur
                MV.write way  j j0
              mj' <- MV.read minv j
              d   <- readSTRef deltaRef
              when (mj' < d) $ do
                writeSTRef deltaRef mj'
                writeSTRef j1Ref j
          delta <- readSTRef deltaRef
          forM_ [0 .. n] $ \j -> do
            isU <- MV.read used j
            if isU
              then do
                pj <- MV.read p j
                upj <- MV.read u pj
                MV.write u pj (upj + delta)
                vj <- MV.read v j
                MV.write v j (vj - delta)
              else do
                mj <- MV.read minv j
                MV.write minv j (mj - delta)
          j1 <- readSTRef j1Ref
          writeSTRef j0Ref j1
          pj1 <- MV.read p j1
          when (pj1 /= 0) step
    step

    -- augmenting path に沿って割当を更新
    let aug = do
          j0 <- readSTRef j0Ref
          j1 <- MV.read way j0
          pj1 <- MV.read p j1
          MV.write p j0 pj1
          writeSTRef j0Ref j1
          when (j1 /= 0) aug
    aug

  -- 結果ベクトルを構築: assignment[i-1] = j-1 (p[j] = i ⇒ row i → col j)
  result <- MV.replicate n (0 :: Int)
  forM_ [1 .. n] $ \j -> do
    pj <- MV.read p j
    when (pj >= 1) $ MV.write result (pj - 1) (j - 1)
  VU.freeze result