competitive/combinatorial_optimization/
maximum_scoring_segment_sets.rs

1use super::{DoublyLinkedList, Zero};
2use std::{
3    cmp::{Ordering, Reverse},
4    ops::{Add, Neg},
5};
6
7/// Returns the maximum total score of exactly `k` nonempty disjoint segments for each `k`.
8pub fn maximum_scoring_segment_sets<T>(scores: &[T]) -> Vec<T>
9where
10    T: Clone + Ord + Zero + Add<Output = T> + Neg<Output = T>,
11{
12    let n = scores.len();
13    let zero = T::zero();
14    let mut positive_sum = zero.clone();
15    let mut nonnegative_count = 0;
16    let mut negative_scores = Vec::with_capacity(n);
17    let mut blocks: Vec<T> = Vec::with_capacity(n);
18
19    for score in scores {
20        let ordering = score.cmp(&zero);
21        match ordering {
22            Ordering::Greater => {
23                positive_sum = positive_sum + score.clone();
24                nonnegative_count += 1;
25            }
26            Ordering::Equal => {
27                nonnegative_count += 1;
28                continue;
29            }
30            Ordering::Less => negative_scores.push(Reverse(score.clone())),
31        }
32
33        if let Some(last) = blocks.last_mut()
34            && ((*last).cmp(&zero) == Ordering::Greater) == (ordering == Ordering::Greater)
35        {
36            *last = last.clone() + score.clone();
37        } else {
38            blocks.push(score.clone());
39        }
40    }
41
42    if blocks.first().is_some_and(|score| score < &zero) {
43        blocks.remove(0);
44    }
45    if blocks.last().is_some_and(|score| score < &zero) {
46        blocks.pop();
47    }
48
49    let mut result = vec![zero.clone(); n + 1];
50    let segment_count = blocks.len().div_ceil(2);
51
52    for value in &mut result[segment_count..=nonnegative_count] {
53        value.clone_from(&positive_sum);
54    }
55
56    negative_scores.sort_unstable();
57    let mut score_with_negatives = positive_sum.clone();
58    for (offset, Reverse(score)) in negative_scores.into_iter().enumerate() {
59        score_with_negatives = score_with_negatives + score;
60        result[nonnegative_count + offset + 1].clone_from(&score_with_negatives);
61    }
62
63    if segment_count <= 1 {
64        return result;
65    }
66
67    let block_count = blocks.len();
68    let mut links = DoublyLinkedList::new(block_count);
69    for index in 1..block_count {
70        links.link(index - 1, index);
71    }
72
73    let absolute = |value: &T| {
74        if value < &zero {
75            -value.clone()
76        } else {
77            value.clone()
78        }
79    };
80    let mut weights: Vec<_> = blocks.iter().map(|block| Some(absolute(block))).collect();
81    let mut candidates: Vec<_> = (0..block_count).collect();
82
83    let mut losses = Vec::with_capacity(segment_count);
84    for _ in 1..segment_count {
85        let (index, left, right) = loop {
86            let index = candidates
87                .pop()
88                .expect("a live block sequence must have a local minimum");
89            if weights[index].is_none() {
90                continue;
91            }
92            let left = links.prev(index);
93            let right = links.next(index);
94            if (left == usize::MAX || weights[index] <= weights[left])
95                && (right == usize::MAX || weights[index] <= weights[right])
96            {
97                break (index, left, right);
98            }
99        };
100        losses.push(weights[index].take().expect("a local minimum must be live"));
101
102        if left == usize::MAX {
103            links.detach(index);
104            let (_, next) = links.detach(right);
105            weights[right] = None;
106            candidates.push(next);
107        } else if right == usize::MAX {
108            let (prev, _) = links.detach(left);
109            links.detach(index);
110            weights[left] = None;
111            candidates.push(prev);
112        } else {
113            blocks[index] = blocks[left].clone() + blocks[index].clone() + blocks[right].clone();
114            weights[index] = Some(absolute(&blocks[index]));
115            let (prev, _) = links.detach(left);
116            let (_, next) = links.detach(right);
117            weights[left] = None;
118            weights[right] = None;
119            candidates.extend(
120                [prev, index, next]
121                    .into_iter()
122                    .filter(|&index| index != usize::MAX),
123            );
124        }
125    }
126
127    losses.extend(weights.into_iter().flatten());
128    losses.sort_unstable();
129
130    let mut current_score = positive_sum;
131    for (result, loss) in result[..segment_count].iter_mut().rev().zip(losses) {
132        current_score = current_score + -loss;
133        result.clone_from(&current_score);
134    }
135    result
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141    use crate::{rand, tools::Xorshift};
142
143    fn naive(scores: &[i32]) -> Vec<i32> {
144        let n = scores.len();
145        let mut result = vec![i32::MIN; n + 1];
146        result[0] = 0;
147        for mask in 1usize..1 << n {
148            let mut score = 0;
149            let mut selected_count = 0;
150            let mut run_count = 0;
151            for (index, &value) in scores.iter().enumerate() {
152                if mask >> index & 1 == 1 {
153                    score += value;
154                    selected_count += 1;
155                    if index == 0 || mask >> (index - 1) & 1 == 0 {
156                        run_count += 1;
157                    }
158                }
159            }
160            for value in &mut result[run_count..=selected_count] {
161                *value = (*value).max(score);
162            }
163        }
164        result
165    }
166
167    #[test]
168    fn test_maximum_scoring_segment_sets() {
169        for n in 0..=6u32 {
170            for mut encoded in 0..5usize.pow(n) {
171                let mut scores = Vec::with_capacity(n as usize);
172                for _ in 0..n {
173                    scores.push((encoded % 5) as i32 - 2);
174                    encoded /= 5;
175                }
176                assert_eq!(maximum_scoring_segment_sets(&scores), naive(&scores));
177            }
178        }
179
180        const Q: usize = 200;
181        const N: usize = 15;
182        const A: i32 = 1000;
183        let mut rng = Xorshift::default();
184        for _ in 0..Q {
185            rand!(rng, n: 7..=N, scores: [-A..=A; n]);
186            assert_eq!(maximum_scoring_segment_sets(&scores), naive(&scores));
187        }
188    }
189}
competitive/combinatorial_optimization/maximum_scoring_segment_sets.rs

competitive/combinatorial_optimization/
maximum_scoring_segment_sets.rs
