{-# LANGUAGE RecordWildCards #-}

-- | Static, \(k\)-dimensional tree \((k = 2)\).
--
-- - Points are fixed on `build`.
-- - Multiple points can exist at the same coordinate.
--
-- ==== __Examples__
-- >>> import AtCoder.Extra.KdTree qualified as KT
-- >>> import Data.Vector.Unboxed qualified as VU
-- >>> let xys = VU.fromList [(0, 0), (1, 1), (4, 2)]
-- >>> let kt = KT.build2 xys
-- >>> -- Find point indices in [0, 2) x [0, 2) with maximum capacity 3
-- >>> KT.findPointsIn kt 0 2 0 2 3
-- [0,1]
--
-- >>> KT.findNearestPoint kt 3 3
-- Just 2
--
-- @since 1.2.2.0
module AtCoder.Extra.KdTree
  ( -- * K-dimensional tree
    KdTree (..),

    -- * Constructors
    build,
    build2,
    findPointsIn,
    findNearestPoint,
  )
where

import AtCoder.Internal.Assert qualified as ACIA
import AtCoder.Internal.Bit qualified as ACIB
import Control.Monad.ST (runST)
import Data.Bits
import Data.Ord (comparing)
import Data.Vector.Algorithms.Intro qualified as VAI
import Data.Vector.Generic qualified as VG
import Data.Vector.Generic.Mutable qualified as VGM
import Data.Vector.Unboxed qualified as VU
import Data.Vector.Unboxed.Mutable qualified as VUM
import GHC.Stack (HasCallStack)

-- | Static, \(k\)-dimensional tree \((k = 2)\).
--
-- @since 1.2.2.0
data KdTree = KdTree
  { -- | The number of points in the \(k\)-d tree.
    --
    -- @since 1.2.2.0
    KdTree -> Int
nKt :: {-# UNPACK #-} !Int,
    -- | Rectangle information: inclusive (closed) ranges \([x_1, x_2) \times [y_1, y_2)\).
    --
    -- @since 1.2.2.0
    KdTree -> Vector (Int, Int, Int, Int)
incRectsKt :: !(VU.Vector (Int, Int, Int, Int)),
    -- | Maps rectangle index to original point index.
    --
    -- @since 1.2.2.0
    KdTree -> Vector Int
dataKt :: !(VU.Vector Int)
  }

-- | \(O(n \log n)\) Creates a `KdTree` from \(x\) and \(y\) vectors.
--
-- ==== Constraints
-- - \(|\mathrm{xs}| = |\mathrm{ys}|\).
--
-- @since 1.2.2.0
{-# INLINEABLE build #-}
build ::
  (HasCallStack) =>
  -- | \(x\) coordnates
  VU.Vector Int ->
  -- | \(y\) coordnates
  VU.Vector Int ->
  -- | `KdTree`
  KdTree
build :: HasCallStack => Vector Int -> Vector Int -> KdTree
build Vector Int
xs0 Vector Int
ys0 =
  let nKt :: Int
nKt = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
xs0
      !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
nKt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
ys0) String
"AtCoder.Extra.KdTree.buildST: the length of `xs`, `ys` and `vs` must be equal"
   in if Int
nKt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        then Int -> Vector (Int, Int, Int, Int) -> Vector Int -> KdTree
KdTree Int
0 Vector (Int, Int, Int, Int)
forall a. Unbox a => Vector a
VU.empty Vector Int
forall a. Unbox a => Vector a
VU.empty
        else (forall s. ST s KdTree) -> KdTree
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s KdTree) -> KdTree)
-> (forall s. ST s KdTree) -> KdTree
forall a b. (a -> b) -> a -> b
$ do
          let vs0 :: Vector Int
vs0 = Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
nKt Int -> Int
forall a. a -> a
id
          let logKt :: Int
logKt = Int -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Int
ACIB.bitCeil (Int
nKt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          MVector (PrimState (ST s)) Int
dat <- Int -> Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int -> Int
forall a. Bits a => Int -> a
bit (Int
logKt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (-Int
1 :: Int)
          MVector (PrimState (ST s)) (Int, Int, Int, Int)
incRectsVec <- Int
-> (Int, Int, Int, Int)
-> ST s (MVector (PrimState (ST s)) (Int, Int, Int, Int))
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> a -> m (MVector (PrimState m) a)
VUM.replicate (Int -> Int
forall a. Bits a => Int -> a
bit (Int
logKt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)) (Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
minBound, Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
minBound)
          let VUM.MV_4 Int
_ MVector (PrimState (ST s)) Int
xMins MVector (PrimState (ST s)) Int
xMaxes MVector (PrimState (ST s)) Int
yMins MVector (PrimState (ST s)) Int
yMaxes = MVector (PrimState (ST s)) (Int, Int, Int, Int)
incRectsVec

          -- - idx: rectangle index (one-based)
          -- - xs, ys, vs: point information (x, y and monoid value)
          -- - ids: maps sorted vertices to the original vertex indices
          -- - divX: represents hyperplane direction for point partition
          let -- buildSubtree :: Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int -> VU.Vector Int -> Bool -> ST s ()
              buildSubtree :: Int
-> Vector Int
-> Vector Int
-> Vector Int
-> Vector a
-> Bool
-> ST s ()
buildSubtree Int
idx Vector Int
xs Vector Int
ys Vector Int
vs Vector a
ids Bool
divX = do
                let n :: Int
n = Vector Int -> Int
forall a. Unbox a => Vector a -> Int
VU.length Vector Int
xs

                -- retrieve the bounds:
                let (!Int
xMin, !Int
xMax, !Int
yMin, !Int
yMax) =
                      ((Int, Int, Int, Int) -> (Int, Int) -> (Int, Int, Int, Int))
-> (Int, Int, Int, Int)
-> Vector (Int, Int)
-> (Int, Int, Int, Int)
forall b a. Unbox b => (a -> b -> a) -> a -> Vector b -> a
VU.foldl'
                        (\(!Int
a, !Int
b, !Int
c, !Int
d) (!Int
x, !Int
y) -> (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
a Int
x, Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
b Int
x, Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
c Int
y, Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
d Int
y))
                        (Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
minBound, Int
forall a. Bounded a => a
maxBound, Int
forall a. Bounded a => a
minBound)
                        (Vector (Int, Int) -> (Int, Int, Int, Int))
-> Vector (Int, Int) -> (Int, Int, Int, Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector (Int, Int)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Int
xs Vector Int
ys
                MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState (ST s)) Int
xMins (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
xMin) Int
idx
                MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState (ST s)) Int
xMaxes (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
xMax) Int
idx
                MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState (ST s)) Int
yMins (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
yMin) Int
idx
                MVector (PrimState (ST s)) Int -> (Int -> Int) -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> (a -> a) -> Int -> m ()
VGM.modify MVector (PrimState (ST s)) Int
yMaxes (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
yMax) Int
idx

                if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1
                  then do
                    -- it's a terminal
                    MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
dat Int
idx (Int -> ST s ()) -> Int -> ST s ()
forall a b. (a -> b) -> a -> b
$ Vector Int
vs Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
0
                  else do
                    -- partition the vertices into two:
                    let m :: Int
m = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
                    let is :: Vector Int
is = (forall s. ST s (MVector s Int)) -> Vector Int
forall a. Unbox a => (forall s. ST s (MVector s a)) -> Vector a
VU.create ((forall s. ST s (MVector s Int)) -> Vector Int)
-> (forall s. ST s (MVector s Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
                          MVector s Int
vec <- Int -> (Int -> Int) -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> (Int -> a) -> m (MVector (PrimState m) a)
VUM.generate Int
n Int -> Int
forall a. a -> a
id
                          if Bool
divX
                            then Comparison Int -> MVector (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
VAI.selectBy ((Int -> Int) -> Comparison Int
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Vector Int
xs VG.!)) MVector s Int
MVector (PrimState (ST s)) Int
vec Int
m
                            else Comparison Int -> MVector (PrimState (ST s)) Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) e.
(PrimMonad m, MVector v e) =>
Comparison e -> v (PrimState m) e -> Int -> m ()
VAI.selectBy ((Int -> Int) -> Comparison Int
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (Vector Int
ys VG.!)) MVector s Int
MVector (PrimState (ST s)) Int
vec Int
m
                          MVector s Int -> ST s (MVector s Int)
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MVector s Int
vec

                    -- TODO: permute in-place?
                    let (!Vector Int
xsL, !Vector Int
xsR) = Int -> Vector Int -> (Vector Int, Vector Int)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
VG.splitAt Int
m (Vector Int -> (Vector Int, Vector Int))
-> Vector Int -> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute Vector Int
xs Vector Int
is
                    let (!Vector Int
ysL, !Vector Int
ysR) = Int -> Vector Int -> (Vector Int, Vector Int)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
VG.splitAt Int
m (Vector Int -> (Vector Int, Vector Int))
-> Vector Int -> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute Vector Int
ys Vector Int
is
                    let (!Vector Int
vsL, !Vector Int
vsR) = Int -> Vector Int -> (Vector Int, Vector Int)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
VG.splitAt Int
m (Vector Int -> (Vector Int, Vector Int))
-> Vector Int -> (Vector Int, Vector Int)
forall a b. (a -> b) -> a -> b
$ Vector Int -> Vector Int -> Vector Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute Vector Int
vs Vector Int
is
                    let (!Vector a
idsL, !Vector a
idsR) = Int -> Vector a -> (Vector a, Vector a)
forall (v :: * -> *) a. Vector v a => Int -> v a -> (v a, v a)
VG.splitAt Int
m (Vector a -> (Vector a, Vector a))
-> Vector a -> (Vector a, Vector a)
forall a b. (a -> b) -> a -> b
$ Vector a -> Vector Int -> Vector a
forall (v :: * -> *) a.
(HasCallStack, Vector v a, Vector v Int) =>
v a -> v Int -> v a
VG.backpermute Vector a
ids Vector Int
is

                    -- build the subtree:
                    Int
-> Vector Int
-> Vector Int
-> Vector Int
-> Vector a
-> Bool
-> ST s ()
buildSubtree (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Vector Int
xsL Vector Int
ysL Vector Int
vsL Vector a
idsL (Bool -> Bool
not Bool
divX)
                    Int
-> Vector Int
-> Vector Int
-> Vector Int
-> Vector a
-> Bool
-> ST s ()
buildSubtree (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Vector Int
xsR Vector Int
ysR Vector Int
vsR Vector a
idsR (Bool -> Bool
not Bool
divX)

          Int
-> Vector Int
-> Vector Int
-> Vector Int
-> Vector Int
-> Bool
-> ST s ()
forall {a}.
(Vector Vector a, MVector MVector a) =>
Int
-> Vector Int
-> Vector Int
-> Vector Int
-> Vector a
-> Bool
-> ST s ()
buildSubtree Int
1 Vector Int
xs0 Vector Int
ys0 Vector Int
vs0 (Int -> (Int -> Int) -> Vector Int
forall a. Unbox a => Int -> (Int -> a) -> Vector a
VU.generate Int
nKt Int -> Int
forall a. a -> a
id) Bool
True
          Vector Int
dataKt <- MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState (ST s)) Int
dat
          Vector (Int, Int, Int, Int)
incRectsKt <- MVector (PrimState (ST s)) (Int, Int, Int, Int)
-> ST s (Vector (Int, Int, Int, Int))
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState (ST s)) (Int, Int, Int, Int)
incRectsVec
          KdTree -> ST s KdTree
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure KdTree {Int
Vector Int
Vector (Int, Int, Int, Int)
nKt :: Int
incRectsKt :: Vector (Int, Int, Int, Int)
dataKt :: Vector Int
nKt :: Int
dataKt :: Vector Int
incRectsKt :: Vector (Int, Int, Int, Int)
..}

-- | \(O(n \log n)\) Creates `KdTree` from a \((x, y)\) vector.
--
-- ==== Constraints
-- - \(|\mathrm{xs}| = |\mathrm{ys}|\).
--
-- @since 1.2.2.0
{-# INLINE build2 #-}
build2 ::
  (HasCallStack) =>
  -- | \(x, y\) coordnates
  VU.Vector (Int, Int) ->
  -- | `KdTree`
  KdTree
build2 :: HasCallStack => Vector (Int, Int) -> KdTree
build2 Vector (Int, Int)
xys = HasCallStack => Vector Int -> Vector Int -> KdTree
Vector Int -> Vector Int -> KdTree
build Vector Int
xs Vector Int
ys
  where
    (!Vector Int
xs, !Vector Int
ys) = Vector (Int, Int) -> (Vector Int, Vector Int)
forall a b.
(Unbox a, Unbox b) =>
Vector (a, b) -> (Vector a, Vector b)
VU.unzip Vector (Int, Int)
xys

-- | \(O(n)\) Collects points in \([x_1, x_2) \times [y_1, y_2)\).
--
-- @since 1.2.2.0
{-# INLINEABLE findPointsIn #-}
findPointsIn ::
  (HasCallStack) =>
  -- | `KdTree`
  KdTree ->
  -- | \(x_1\)
  Int ->
  -- | \(x_2\)
  Int ->
  -- | \(y_1\)
  Int ->
  -- | \(y_2\)
  Int ->
  -- | Maximum number of points in \([x_1, x_2) \times [y_1, y_2)\).
  Int ->
  -- | Point indices in \([x_1, x_2) \times [y_1, y_2)\).
  VU.Vector Int
findPointsIn :: HasCallStack =>
KdTree -> Int -> Int -> Int -> Int -> Int -> Vector Int
findPointsIn KdTree {Int
Vector Int
Vector (Int, Int, Int, Int)
nKt :: KdTree -> Int
incRectsKt :: KdTree -> Vector (Int, Int, Int, Int)
dataKt :: KdTree -> Vector Int
nKt :: Int
incRectsKt :: Vector (Int, Int, Int, Int)
dataKt :: Vector Int
..} Int
x1 Int
x2 Int
y1 Int
y2 Int
capacity
  | Int
nKt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Vector Int
forall a. Unbox a => Vector a
VU.empty
  | Bool
otherwise = (forall s. ST s (Vector Int)) -> Vector Int
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vector Int)) -> Vector Int)
-> (forall s. ST s (Vector Int)) -> Vector Int
forall a b. (a -> b) -> a -> b
$ do
      MVector (PrimState (ST s)) Int
res <- Int -> ST s (MVector (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Unbox a) =>
Int -> m (MVector (PrimState m) a)
VUM.unsafeNew (Int -> ST s (MVector (PrimState (ST s)) Int))
-> Int -> ST s (MVector (PrimState (ST s)) Int)
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
nKt Int
capacity
      let inner :: Int -> Int -> ST s Int
inner Int
i Int
iPush
            -- not intersected
            | Int
x2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
xMin Bool -> Bool -> Bool
|| Int
xMax Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
x1 = Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iPush
            | Int
y2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
yMin Bool -> Bool -> Bool
|| Int
yMax Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
y1 = Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
iPush
            -- a leaf
            | Int
vi Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= -Int
1 = do
                MVector (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall (m :: * -> *) (v :: * -> * -> *) a.
(HasCallStack, PrimMonad m, MVector v a) =>
v (PrimState m) a -> Int -> a -> m ()
VGM.write MVector (PrimState (ST s)) Int
res Int
iPush Int
vi
                Int -> ST s Int
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int -> ST s Int) -> Int -> ST s Int
forall a b. (a -> b) -> a -> b
$ Int
iPush Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            -- a parental rectangle area
            | Bool
otherwise = do
                Int
iPush' <- Int -> Int -> ST s Int
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) Int
iPush
                Int -> Int -> ST s Int
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
iPush'
            where
              (!Int
xMin, !Int
xMax, !Int
yMin, !Int
yMax) = Vector (Int, Int, Int, Int)
incRectsKt Vector (Int, Int, Int, Int) -> Int -> (Int, Int, Int, Int)
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
              vi :: Int
vi = Vector Int
dataKt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
      Int
n <- Int -> Int -> ST s Int
inner Int
1 Int
0
      Int -> Vector Int -> Vector Int
forall a. Unbox a => Int -> Vector a -> Vector a
VU.take Int
n (Vector Int -> Vector Int)
-> ST s (Vector Int) -> ST s (Vector Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVector (PrimState (ST s)) Int -> ST s (Vector Int)
forall a (m :: * -> *).
(Unbox a, PrimMonad m) =>
MVector (PrimState m) a -> m (Vector a)
VU.unsafeFreeze MVector (PrimState (ST s)) Int
res
  where
    !()
_ = HasCallStack => Bool -> String -> ()
Bool -> String -> ()
ACIA.runtimeAssert (Int
x1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
x2 Bool -> Bool -> Bool
&& Int
y1 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
y2) String
"AtCoder.Extra.KdTree.findPointsIn: given invalid interval"

-- | \(O(\log n)\), only if the points are randomly distributed. Returns the index of the nearest
-- point, or `Nothing` if the `KdTree` has no point.
--
-- @since 1.2.2.0
{-# INLINEABLE findNearestPoint #-}
findNearestPoint ::
  (HasCallStack) =>
  -- | `KdTree`
  KdTree ->
  -- | \(x\)
  Int ->
  -- | \(y\)
  Int ->
  -- | The nearest point index
  Maybe Int
findNearestPoint :: HasCallStack => KdTree -> Int -> Int -> Maybe Int
findNearestPoint KdTree {Int
Vector Int
Vector (Int, Int, Int, Int)
nKt :: KdTree -> Int
incRectsKt :: KdTree -> Vector (Int, Int, Int, Int)
dataKt :: KdTree -> Vector Int
nKt :: Int
incRectsKt :: Vector (Int, Int, Int, Int)
dataKt :: Vector Int
..} Int
x Int
y
  | Int
nKt Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Maybe Int
forall a. Maybe a
Nothing
  | Bool
otherwise = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int -> Maybe Int)
-> ((Int, Int) -> Int) -> (Int, Int) -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> Int
forall a b. (a, b) -> a
fst ((Int, Int) -> Maybe Int) -> (Int, Int) -> Maybe Int
forall a b. (a -> b) -> a -> b
$! Int -> (Int, Int) -> (Int, Int)
inner Int
1 {- FIXME: -} (-Int
1, -Int
1)
  where
    clamp :: a -> a -> a -> a
clamp a
a a
aMin a
aMax = a -> a -> a
forall a. Ord a => a -> a -> a
min a
aMax (a -> a) -> a -> a
forall a b. (a -> b) -> a -> b
$ a -> a -> a
forall a. Ord a => a -> a -> a
max a
a a
aMin
    -- Used for pruning. It's |(x, y)|^2 if the (x, y) is within the rectangle.
    bestDistSquared :: Int -> Int
bestDistSquared Int
i =
      let (!Int
xMin, !Int
xMax, !Int
yMin, !Int
yMax) = Vector (Int, Int, Int, Int)
incRectsKt Vector (Int, Int, Int, Int) -> Int -> (Int, Int, Int, Int)
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
          dx :: Int
dx = Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int -> Int -> Int
forall {a}. Ord a => a -> a -> a -> a
clamp Int
x Int
xMin Int
xMax
          dy :: Int
dy = Int
y Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int -> Int -> Int -> Int
forall {a}. Ord a => a -> a -> a -> a
clamp Int
y Int
yMin Int
yMax
       in Int
dx Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
dx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
dy Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
dy
    -- returns (index, bestDist)
    inner :: Int -> (Int, Int) -> (Int, Int)
inner Int
i res :: (Int, Int)
res@(!Int
resV, !Int
resD)
      -- pruning (we have a better point than any point in this rectangle)
      | Int
resV Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= -Int
1 Bool -> Bool -> Bool
&& Int
resD Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
d = (Int, Int)
res
      -- it's a leaf
      | Int
dataI Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= -Int
1 = (Int
dataI, Int
d)
      -- look into the children
      | Int
d0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
d1 = Int -> (Int, Int) -> (Int, Int)
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) ((Int, Int) -> (Int, Int)) -> (Int, Int) -> (Int, Int)
forall a b. (a -> b) -> a -> b
$ Int -> (Int, Int) -> (Int, Int)
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int, Int)
res
      | Bool
otherwise = Int -> (Int, Int) -> (Int, Int)
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int, Int) -> (Int, Int)) -> (Int, Int) -> (Int, Int)
forall a b. (a -> b) -> a -> b
$ Int -> (Int, Int) -> (Int, Int)
inner (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0) (Int, Int)
res
      where
        d :: Int
d = Int -> Int
bestDistSquared Int
i
        dataI :: Int
dataI = Vector Int
dataKt Vector Int -> Int -> Int
forall (v :: * -> *) a.
(HasCallStack, Vector v a) =>
v a -> Int -> a
VG.! Int
i
        d0 :: Int
d0 = Int -> Int
bestDistSquared (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0)
        d1 :: Int
d1 = Int -> Int
bestDistSquared (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
0)