Kth Sum

[tipstar0125]

• ABC281E: https://atcoder.jp/contests/abc281/tasks/abc281_e
• ABC306E: https://atcoder.jp/contests/abc306/tasks/abc306_e

[tipstar0125]

• 昇順Ver

#[derive(Debug, Clone)]
struct KthSum {
    N: usize,
    K: usize,
    S: isize,
    large: BTreeSet<(isize, usize)>,
    small: BTreeSet<(isize, usize)>,
}

impl KthSum {
    fn new(N: usize, K: usize) -> Self {
        let mut large = BTreeSet::new();
        let mut small = BTreeSet::new();
        for i in 0..K {
            small.insert((0, i));
        }
        for i in K..N {
            large.insert((0, i));
        }
        KthSum {
            N,
            K,
            S: 0,
            large,
            small,
        }
    }
    fn balance(&mut self) {
        while self.small.len() > self.K {
            let p = *self.small.iter().next_back().unwrap();
            self.small.remove(&p);
            self.large.insert(p);
            self.S -= p.0;
        }
        if self.small.is_empty() || self.large.is_empty() {
            return;
        }
        while self.large.iter().next().unwrap().0 < self.small.iter().next_back().unwrap().0 {
            let large_min = *self.large.iter().next().unwrap();
            let small_max = *self.small.iter().next_back().unwrap();
            self.large.remove(&large_min);
            self.small.remove(&small_max);
            self.large.insert(small_max);
            self.small.insert(large_min);
            self.S += large_min.0 - small_max.0;
        }
    }
    fn add(&mut self, p: (isize, usize)) {
        self.small.insert(p);
        self.S += p.0;
        self.balance();
    }
    fn remove(&mut self, p: (isize, usize)) {
        if self.small.remove(&p) {
            self.S -= p.0;
        } else {
            self.large.remove(&p);
        }
        self.balance();
    }
}

[tipstar0125]

• 降順Ver

#[derive(Debug, Clone)]
struct KthSum {
    N: usize,
    K: usize,
    S: isize,
    large: BTreeSet<(isize, usize)>,
    small: BTreeSet<(isize, usize)>,
}

impl KthSum {
    fn new(N: usize, K: usize) -> Self {
        let mut large = BTreeSet::new();
        let mut small = BTreeSet::new();
        for i in 0..K {
            large.insert((0, i));
        }
        for i in K..N {
            small.insert((0, i));
        }
        KthSum {
            N,
            K,
            S: 0,
            large,
            small,
        }
    }
    fn balance(&mut self) {
        while self.large.len() > self.K {
            let p = *self.large.iter().next().unwrap();
            self.large.remove(&p);
            self.small.insert(p);
            self.S -= p.0;
        }
        if self.small.is_empty() || self.large.is_empty() {
            return;
        }
        while self.large.iter().next().unwrap().0 < self.small.iter().next_back().unwrap().0 {
            let large_min = *self.large.iter().next().unwrap();
            let small_max = *self.small.iter().next_back().unwrap();
            self.large.remove(&large_min);
            self.small.remove(&small_max);
            self.large.insert(small_max);
            self.small.insert(large_min);
            self.S += small_max.0 - large_min.0;
        }
    }
    fn add(&mut self, p: (isize, usize)) {
        self.large.insert(p);
        self.S += p.0;
        self.balance();
    }
    fn remove(&mut self, p: (isize, usize)) {
        if self.large.remove(&p) {
            self.S -= p.0;
        } else {
            self.small.remove(&p);
        }
        self.balance();
    }
}

[tipstar0125]

• add: Kthに入る方（昇順の場合はsmall、降順の場合はlarge）に追加して、Sumを更新し、バランスさせる
• remove: small or largeの存在する方から除去して、Kthの方から除去した場合は、Sumを更新し、バランスさせる
• balnace:
  • Kthの方の長さがKを超えている場合は、small側に移動させる
  • largeのmin < smallのmaxならば、スワップして、Sumを差分だけ更新

[tipstar0125]

Rustはmultisetがないので、タプルで重複しないユニークな値（例えば、配列のindex）を使うことで同様の処理を行う。 参考：https://namacha411.hatenablog.com/entry/2022/02/28/194540