cp_library_rs/algebraic_structure/
k_best.rs

1//! 上位 K 個を保持するモノイド
2
3#![allow(clippy::type_complexity)]
4
5use std::cmp::Ordering;
6
7use crate::algebraic_structure::monoid_with_context::MonoidCtx;
8
9/// 上位 K 個を保持するモノイド
10pub struct KBest<const K: usize, T> {
11    /// 比較関数: ge(a, b) := a ≥ b
12    ge: Box<dyn Fn(&T, &T) -> bool>,
13}
14
15impl<const K: usize, T> KBest<K, T> {
16    /// 比較関数 ge を受取り,モノイド KBest を返す.
17    /// - 比較関数: `ge(a, b) := a ≥ b`
18    /// - ge は推移律を満たす必要がある
19    pub fn new(ge: impl Fn(&T, &T) -> bool + 'static) -> Self {
20        Self { ge: Box::new(ge) }
21    }
22}
23
24impl<const K: usize, T: Clone> MonoidCtx for KBest<K, T> {
25    type Val = [Option<T>; K];
26    fn e(&self) -> Self::Val {
27        std::array::from_fn(|_| None)
28    }
29    fn op(&self, left: &Self::Val, right: &Self::Val) -> Self::Val {
30        let mut res = Vec::with_capacity(K * 2);
31        res.extend_from_slice(left);
32        res.extend_from_slice(right);
33        res.sort_unstable_by(|a, b| match (a, b) {
34            (None, None) => Ordering::Equal,
35            (None, Some(_)) => Ordering::Greater,
36            (Some(_), None) => Ordering::Less,
37            (Some(x), Some(y)) => {
38                let xy = (self.ge)(x, y);
39                let yx = (self.ge)(y, x);
40                if xy && yx {
41                    Ordering::Equal
42                } else if xy {
43                    Ordering::Less
44                } else {
45                    Ordering::Greater
46                }
47            }
48        });
49        std::array::from_fn(|i| res[i].clone())
50    }
51}