maximum-segment-sum-after-removals.rs

use std::cmp::*;

#[derive(Debug, Clone, PartialOrd, PartialEq, Eq)]
struct Zero {
    index: usize,
}

#[derive(Debug, Clone, PartialOrd, PartialEq, Eq, Ord)]
struct Segment {
    start: usize,
    end: usize,
}

#[derive(Debug, Clone)]
enum Node {
    Zero(Zero),
    Segment(Segment),
}

impl Segment {
    fn sum(&self, arr: &[i32]) -> i64 {
        arr[self.start..=self.end].iter().map(|c| {println!("{:?}", *c); *c as i64}).sum()
    }
    fn new_from_slice(arr: &[i32]) -> Self {
        Segment {
            start: 0,
            end: arr.len() - 1,
        }
    }
    fn split_at(self, old_sum: i64, pos: usize, arr: &[i32]) -> (Option<(Self, i64)>, Option<(Self, i64)>) {
        let left = 
            if pos > self.start {
                let seg = Segment {
                    start: self.start,
                    end: pos - 1,
                };
                let sum = seg.sum(arr);
                Some((seg, sum))
            }else {
                    None
            };
        let right = 
            if pos < self.end {
                let seg = Segment {
                    start: pos + 1,
                    end: self.end,
                };
                let sum = seg.sum(arr);
                Some((seg, sum))
            }else {
                    None
            };
        (left, right)
    }
}

impl Node {
    pub fn new(segment: Segment) -> Self {
        Node::Segment(segment)
    }
    pub fn zero(pos: i32) -> Self {
        Node::Zero(Zero {
            index: pos as usize,
        })
    }
}

impl Ord for Node {
    fn cmp(&self, other: &Self) -> Ordering {
        match (self, other) {
            (Node::Segment(s), Node::Segment(o)) => {
                if s.start == o.start && s.start == o.end {
                    Ordering::Equal
                } else if s.end < o.start && s.start < o.start {
                    Ordering::Less
                } else if s.start > o.end && s.start > o.start {
                    Ordering::Greater
                } else {
                    panic!("Bad ordering or overlapping segments")
                }
            }
            (Node::Zero(s), Node::Zero(o)) => s.index.cmp(&o.index),
            (Node::Zero(s), Node::Segment(o)) => {
                if s.index < o.start {
                    Ordering::Less
                } else if s.index > o.end {
                    Ordering::Greater
                } else if (o.start..=o.end).contains(&s.index) {
                    Ordering::Equal
                } else {
                    panic!("Can't compare")
                }
            }
            (Node::Segment(s), Node::Zero(o)) => {
                if o.index < s.start {
                    Ordering::Greater
                } else if o.index > s.end {
                    Ordering::Less
                } else if (s.start..=s.end).contains(&o.index) {
                    Ordering::Equal
                } else {
                    panic!("Can't compare")
                }
            }
        }
    }
    fn max(self, other: Self) -> Self {
        unimplemented!()
    }
    fn min(self, other: Self) -> Self {
        unimplemented!()
    }
    fn clamp(self, min: Self, max: Self) -> Self {
        unimplemented!()
    }
}

impl std::cmp::PartialOrd for Node {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for Node {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Node::Zero(s), Node::Zero(o)) => s.index.eq(&o.index),
            (Node::Segment(s), Node::Segment(o)) => s.start.eq(&o.start) && s.end.eq(&o.end),
            (Node::Segment(s), Node::Zero(o)) => (s.start..=s.end).contains(&o.index),
            (Node::Zero(s), Node::Segment(o)) => (o.start..=o.end).contains(&s.index),
        }
    }
}

impl Eq for Node {}

struct Solution {}
impl Solution {
    pub fn maximum_segment_sum(nums: Vec<i32>, remove_queries: Vec<i32>) -> Vec<i64> {
        let mut output = Vec::with_capacity(remove_queries.len());
        let mut segments = std::collections::BTreeMap::new();
        let mut sums = std::collections::BTreeSet::new();
        let root_segment = Segment::new_from_slice(&nums);
        let root_sum = root_segment.sum(&nums);
        segments.insert(Node::new(root_segment), root_sum);
        sums.insert(root_sum);
        println!("{:?}, {:?}", nums, sums);
        for remove_index in remove_queries.iter() {
            match segments.remove_entry(&Node::zero(*remove_index)) {
                None => panic!("remove outside of range"),
                Some((Node::Zero(_), _)) => panic!("Duplicate removal!"), // I can do it, but why?
                Some((Node::Segment(seg), seg_sum)) => {
                    sums.remove(&seg_sum);
                    let (left, right) =
                        seg.split_at(seg_sum, *remove_index as usize, &nums);
                        if let Some((seg, sum)) = left{
                            segments.insert(Node::new(seg), sum);
                            sums.insert(sum);
                        }
                        if let Some((seg, sum)) = right{
                            segments.insert(Node::new(seg), sum);
                            sums.insert(sum);
                        }
                    output.push(*sums.last().unwrap());
                }
                _ => unimplemented!(),
            }
        }
        output
    }
}

fn main() {
    assert_eq!(
        Solution::maximum_segment_sum(vec![3, 2, 11, 1], vec![3, 2, 1, 0]),
        vec![16, 5, 3, 0]
    );
    assert_eq!(
        Solution::maximum_segment_sum(vec![1, 2, 5, 6, 1], vec![0, 3, 2, 4, 1]),
        vec![14, 7, 2, 2, 0]
    );
}