competitive/combinatorial_optimization/
lexicographical_subsequence.rs

1#[derive(Debug, Clone)]
2pub struct LexicographicalSubsequence {
3    dp: Vec<usize>,
4    index: Vec<Vec<usize>>,
5}
6
7impl LexicographicalSubsequence {
8    pub fn new(sequence: &[usize]) -> Self {
9        let n = sequence.len();
10        let w = sequence.iter().max().map(|w| w + 1).unwrap_or_default();
11        let mut dp = vec![0usize; n + 2];
12        dp[n] = 1;
13        let mut next = vec![n; w];
14        let mut index = vec![vec![]; w];
15        for (i, c) in sequence.iter().cloned().enumerate().rev() {
16            next[c] = i;
17            index[c].push(i);
18            dp[i] = next.iter().fold(1, |acc, &j| acc.saturating_add(dp[j + 1]));
19        }
20        Self { dp, index }
21    }
22
23    /// empty sequence is included
24    pub fn kth_sequence(&self, mut k: usize) -> Option<Vec<usize>> {
25        if self.dp[0] <= k {
26            return None;
27        }
28        let mut seq = Vec::new();
29        let mut pos: Vec<_> = self.index.iter().map(Vec::len).collect();
30        let mut cur = 0usize;
31        while k > 0 {
32            k -= 1;
33            for (c, (idx, pos)) in self.index.iter().zip(pos.iter_mut()).enumerate() {
34                while *pos > 0 && idx[*pos - 1] < cur {
35                    *pos -= 1;
36                }
37                if *pos == 0 {
38                    continue;
39                }
40                let x = idx[*pos - 1];
41                if self.dp[x + 1] <= k {
42                    k -= self.dp[x + 1];
43                } else {
44                    cur = x + 1;
45                    seq.push(c);
46                    break;
47                }
48            }
49        }
50        Some(seq)
51    }
52}