webgraph_algo/utils/
argmax.rs

1/*
2 * SPDX-FileCopyrightText: 2024 Matteo Dell'Acqua
3 * SPDX-FileCopyrightText: 2025 Sebastiano Vigna
4 *
5 * SPDX-License-Identifier: Apache-2.0 OR LGPL-2.1-or-later
6 */
7
8/// Returns the index of the maximum value in an iterator, or [`None`] if the
9/// iterator is empty.
10///
11/// If the maximum appears several times, this methods returns the position of
12/// the first instance.
13///
14/// # Arguments
15///
16/// * `iter`: the iterator.
17///
18/// # Panics
19///
20/// If a comparison returns [`None`].
21///
22///
23/// # Examples
24/// ```
25/// # use webgraph_algo::utils::math::argmax;
26/// let v = vec![1, 2, 5, 2, 1, 5];
27/// let index = argmax(&v);
28/// assert_eq!(index, Some(2));
29/// ```
30pub fn argmax<I: IntoIterator>(iter: I) -> Option<usize>
31where
32    I::Item: core::cmp::PartialOrd + Copy,
33{
34    iter.into_iter()
35        .enumerate()
36        .min_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap())
37        .map(|(idx, _)| idx)
38}
39
40/// Returns the index of the maximum value approved by a filter in an iterator,
41/// or [`None`] if no element is approved by the filter.
42///
43/// In case of ties, this method returns the index for which the corresponding
44/// element in `tie_break` is minimized.
45///
46/// If the maximum appears several times with the same tie break, this methods
47/// returns the position of the first instance.
48///
49/// # Arguments
50///
51/// * `iter`: the iterator.
52///
53/// * `tie_break`: in case two elements of `iter` are the same, the
54///   corresponding elements in this iterator are used as secondary order.
55///
56/// * `filter`: a closure that takes as arguments the index of the element and
57///   the element itself and returns true if the element is approved.
58///
59/// # Panics
60///
61/// If a comparison returns [`None`].
62///
63/// # Examples
64///
65/// ```rust
66/// # use webgraph_algo::utils::math::argmax_filtered;
67/// let v = vec![1, 2, 5, 2, 1, 2];
68/// let tie = vec![1, 2, 3, 4, 5, 2];
69/// let index = argmax_filtered(&v, &tie, |_, &element| element < 4);
70/// // Tie break wins
71/// assert_eq!(index, Some(3));
72///
73/// let v = vec![1, 2, 5, 2, 1, 2];
74/// let tie = vec![1, 1, 3, 2, 5, 2];
75/// let index = argmax_filtered(&v, &tie, |_, &element| element < 4);
76/// // Enumeration order wins
77/// assert_eq!(index, Some(3));
78/// ```
79pub fn argmax_filtered<I: IntoIterator, J: IntoIterator>(
80    iter: I,
81    tie_break: J,
82    filter: impl Fn(usize, I::Item) -> bool,
83) -> Option<usize>
84where
85    I::Item: core::cmp::PartialOrd + Copy,
86    J::Item: core::cmp::PartialOrd + Copy,
87{
88    iter.into_iter()
89        .zip(tie_break)
90        .enumerate()
91        .filter(|(idx, (v, _tie))| filter(*idx, *v))
92        .min_by(|(_, a), (_, b)| b.partial_cmp(a).unwrap())
93        .map(|(idx, _)| idx)
94}