competitive/algorithm/
ternary_search.rs

1use std::ops::RangeInclusive;
2
3/// fibonacci search helper
4pub trait FibonacciSearch: Sized {
5    fn fibonacci_search<T, F>(self, other: Self, f: F) -> (Self, T)
6    where
7        T: PartialOrd,
8        F: FnMut(Self) -> T;
9}
10macro_rules! impl_fibonacci_search_unsigned {
11    ($($t:ty)*) => {
12        $(impl FibonacciSearch for $t {
13            fn fibonacci_search<T, F>(self, other: Self, mut f: F) -> (Self, T)
14            where
15                T: PartialOrd,
16                F: FnMut(Self) -> T,
17            {
18                let l = self;
19                let r = other;
20                assert!(l <= r);
21                const W: usize = [12, 23, 46, 92, 185][<$t>::BITS.ilog2() as usize - 3];
22                const FIB: [$t; W] = {
23                    let mut fib = [0; W];
24                    fib[0] = 1;
25                    fib[1] = 2;
26                    let mut i = 2;
27                    while i < W {
28                        fib[i] = fib[i - 1] + fib[i - 2];
29                        i += 1;
30                    }
31                    fib
32                };
33                let mut s = l;
34                let mut v0 = None;
35                let mut v1 = None;
36                let mut v2 = None;
37                let mut v3 = None;
38                for w in FIB[..FIB.partition_point(|&f| f < r - l)].windows(2).rev() {
39                    let (w0, w1) = (w[0], w[1]);
40                    if w1 > r - s || v1.get_or_insert_with(|| f(s + w0)) <= v2.get_or_insert_with(|| f(s + w1)) {
41                        v3 = v2;
42                        v2 = v1;
43                        v1 = None;
44                    } else {
45                        v0 = v1;
46                        v1 = v2;
47                        v2 = None;
48                        s += w0;
49                    }
50                }
51                let mut kv = (s, v0.unwrap_or_else(|| f(s)));
52                if s < r {
53                    let v = v1.or(v2).unwrap_or_else(|| f(s + 1));
54                    if v < kv.1 {
55                        kv = (s + 1, v);
56                    }
57                    if s + 1 < r {
58                        let v = v3.unwrap_or_else(|| f(s + 2));
59                        if v < kv.1 {
60                            kv = (s + 2, v);
61                        }
62                    }
63                }
64                kv
65            }
66        })*
67    };
68}
69impl_fibonacci_search_unsigned!(u8 u16 u32 u64 u128 usize);
70
71/// ternary search helper
72pub trait Trisect: Clone {
73    type Key: FibonacciSearch;
74    fn trisect_key(self) -> Self::Key;
75    fn trisect_unkey(key: Self::Key) -> Self;
76}
77
78macro_rules! impl_trisect_unsigned {
79    ($($t:ty)*) => {
80        $(impl Trisect for $t {
81            type Key = $t;
82            fn trisect_key(self) -> Self::Key {
83                self
84            }
85            fn trisect_unkey(key: Self::Key) -> Self {
86                key
87            }
88        })*
89    };
90}
91macro_rules! impl_trisect_signed {
92    ($({$i:ident $u:ident})*) => {
93        $(impl Trisect for $i {
94            type Key = $u;
95            fn trisect_key(self) -> Self::Key {
96                (self as $u) ^ (1 << <$u>::BITS - 1)
97            }
98            fn trisect_unkey(key: Self::Key) -> Self {
99                (key ^ (1 << <$u>::BITS - 1)) as $i
100            }
101        })*
102    };
103}
104macro_rules! impl_trisect_float {
105    ($({$t:ident $u:ident $i:ident})*) => {
106        $(impl Trisect for $t {
107            type Key = $u;
108            fn trisect_key(self) -> Self::Key {
109                let a = self.to_bits() as $i;
110                (a ^ (((a >> <$u>::BITS - 1) as $u) >> 1) as $i) as $u ^ (1 << <$u>::BITS - 1)
111            }
112            fn trisect_unkey(key: Self::Key) -> Self {
113                let key = (key  ^ (1 << <$u>::BITS - 1)) as $i;
114                $t::from_bits((key ^ (((key >> <$u>::BITS - 1) as $u) >> 1) as $i) as _)
115            }
116        })*
117    };
118}
119
120impl_trisect_unsigned!(u8 u16 u32 u64 u128 usize);
121impl_trisect_signed!({i8 u8} {i16 u16} {i32 u32} {i64 u64} {i128 u128} {isize usize});
122impl_trisect_float!({f32 u32 i32} {f64 u64 i64});
123
124/// Returns the element that gives the minimum value from the strictly concave up function and the minimum value.
125pub fn ternary_search<K, V, F>(range: RangeInclusive<K>, mut f: F) -> (K, V)
126where
127    K: Trisect,
128    V: PartialOrd,
129    F: FnMut(K) -> V,
130{
131    let (l, r) = range.into_inner();
132    let (k, v) =
133        <K::Key as FibonacciSearch>::fibonacci_search(l.trisect_key(), r.trisect_key(), |x| {
134            f(Trisect::trisect_unkey(x))
135        });
136    (K::trisect_unkey(k), v)
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use crate::num::DoubleDouble;
143
144    #[test]
145    fn test_ternary_search_unsigned() {
146        for p in 0u8..=u8::MAX {
147            for l in 0u8..=u8::MAX {
148                for r in l..=u8::MAX {
149                    let f = |x| p.abs_diff(x);
150                    assert_eq!(
151                        f(l).min(f(r)).min(f(p.clamp(l, r))),
152                        ternary_search(l..=r, f).1,
153                    );
154                }
155            }
156        }
157    }
158
159    #[test]
160    fn test_ternary_search_signed() {
161        for p in -20..=20 {
162            assert_eq!(
163                p.clamp(-10, 10),
164                ternary_search(-10i64..=10, |x| 10 * (x - p).pow(2) + 5).0,
165            );
166        }
167    }
168
169    #[test]
170    fn test_ternary_search_float() {
171        assert_eq!(
172            std::f64::consts::PI,
173            ternary_search(f64::MIN..=f64::MAX, |x| (DoubleDouble::from(x)
174                - DoubleDouble::from(std::f64::consts::PI))
175            .abs())
176            .0,
177        );
178    }
179}
competitive/algorithm/ternary_search.rs

competitive/algorithm/
ternary_search.rs
