Finding mk average

use std::collections::{BTreeMap, VecDeque};

pub struct MKAverage {
  m: usize,
  k: usize,
  window: VecDeque<i32>,
  // 3 sorted partitions: low (k smallest), mid (middle m-2k), high (k largest)
  low: BTreeMap<i32, i32>,
  mid: BTreeMap<i32, i32>,
  high: BTreeMap<i32, i32>,
  low_cnt: usize,
  mid_cnt: usize,
  high_cnt: usize,
  mid_sum: i64,
}

impl MKAverage {
  fn map_inc(m: &mut BTreeMap<i32, i32>, v: i32) { *m.entry(v).or_insert(0) += 1; }
  fn map_dec(m: &mut BTreeMap<i32, i32>, v: i32) {
    if let Some(c) = m.get_mut(&v) { *c -= 1; if *c == 0 { m.remove(&v); } }
  }
  fn map_max(m: &BTreeMap<i32, i32>) -> i32 { *m.keys().next_back().unwrap() }
  fn map_min(m: &BTreeMap<i32, i32>) -> i32 { *m.keys().next().unwrap() }

  // Insert val into low, then cascade to restore invariant
  fn insert_cascade(&mut self, val: i32) {
    Self::map_inc(&mut self.low, val); self.low_cnt += 1;
    if self.low_cnt > self.k {
      let mv = Self::map_max(&self.low);
      Self::map_dec(&mut self.low, mv); self.low_cnt -= 1;
      Self::map_inc(&mut self.mid, mv); self.mid_cnt += 1; self.mid_sum += mv as i64;
    }
    let target_mid = self.m - 2 * self.k;
    if self.mid_cnt > target_mid {
      let mv = Self::map_max(&self.mid);
      Self::map_dec(&mut self.mid, mv); self.mid_cnt -= 1; self.mid_sum -= mv as i64;
      Self::map_inc(&mut self.high, mv); self.high_cnt += 1;
    }
    // Fix ordering: if max(low) > min(mid), swap them
    if self.low_cnt > 0 && self.mid_cnt > 0 {
      let lmax = Self::map_max(&self.low);
      let mmin = Self::map_min(&self.mid);
      if lmax > mmin {
        Self::map_dec(&mut self.low, lmax); Self::map_inc(&mut self.mid, lmax); self.mid_sum += lmax as i64;
        Self::map_dec(&mut self.mid, mmin); self.mid_sum -= mmin as i64; Self::map_inc(&mut self.low, mmin);
      }
    }
    // Fix ordering: if max(mid) > min(high), swap them
    if self.mid_cnt > 0 && self.high_cnt > 0 {
      let mmax = Self::map_max(&self.mid);
      let hmin = Self::map_min(&self.high);
      if mmax > hmin {
        Self::map_dec(&mut self.mid, mmax); self.mid_sum -= mmax as i64; Self::map_inc(&mut self.high, mmax);
        Self::map_dec(&mut self.high, hmin); Self::map_inc(&mut self.mid, hmin); self.mid_sum += hmin as i64;
      }
    }
  }

  // Remove val, then rebalance
  fn remove_rebalance(&mut self, val: i32) {
    // Find which partition contains val
    let in_low = self.low_cnt > 0 && val <= Self::map_max(&self.low);
    let in_high = self.high_cnt > 0 && val >= Self::map_min(&self.high);
    if in_low {
      Self::map_dec(&mut self.low, val); self.low_cnt -= 1;
      // Pull min of mid into low
      if self.low_cnt < self.k && self.mid_cnt > 0 {
        let mv = Self::map_min(&self.mid);
        Self::map_dec(&mut self.mid, mv); self.mid_cnt -= 1; self.mid_sum -= mv as i64;
        Self::map_inc(&mut self.low, mv); self.low_cnt += 1;
      }
      // If mid is now short, pull min of high
      let target_mid = self.m - 2 * self.k;
      if self.mid_cnt < target_mid && self.high_cnt > 0 {
        let mv = Self::map_min(&self.high);
        Self::map_dec(&mut self.high, mv); self.high_cnt -= 1;
        Self::map_inc(&mut self.mid, mv); self.mid_cnt += 1; self.mid_sum += mv as i64;
      }
    } else if in_high {
      Self::map_dec(&mut self.high, val); self.high_cnt -= 1;
      // Pull max of mid into high
      if self.high_cnt < self.k && self.mid_cnt > 0 {
        let mv = Self::map_max(&self.mid);
        Self::map_dec(&mut self.mid, mv); self.mid_cnt -= 1; self.mid_sum -= mv as i64;
        Self::map_inc(&mut self.high, mv); self.high_cnt += 1;
      }
      // If mid is now short, pull max of low
      let target_mid = self.m - 2 * self.k;
      if self.mid_cnt < target_mid && self.low_cnt > 0 {
        let mv = Self::map_max(&self.low);
        Self::map_dec(&mut self.low, mv); self.low_cnt -= 1;
        Self::map_inc(&mut self.mid, mv); self.mid_cnt += 1; self.mid_sum += mv as i64;
      }
    } else {
      Self::map_dec(&mut self.mid, val); self.mid_cnt -= 1; self.mid_sum -= val as i64;
      // Mid lost one. Pull from high if possible (prefer high since we also check low)
      if self.high_cnt > 0 {
        let mv = Self::map_min(&self.high);
        Self::map_dec(&mut self.high, mv); self.high_cnt -= 1;
        Self::map_inc(&mut self.mid, mv); self.mid_cnt += 1; self.mid_sum += mv as i64;
        // High is short, pull from mid's new max? No - just pull max of low to high if low is over k
        // Actually high is now k-1, we need to check later
        // For simplicity: pull max of low to mid then push max of mid to high
        // But that might cause cascade. Let me just ensure high stays at k:
        if self.low_cnt > self.k {
          let mv2 = Self::map_max(&self.low);
          Self::map_dec(&mut self.low, mv2); self.low_cnt -= 1;
          Self::map_inc(&mut self.high, mv2); self.high_cnt += 1;
        }
      } else if self.low_cnt > 0 {
        // Pull max of low to mid
        let mv = Self::map_max(&self.low);
        Self::map_dec(&mut self.low, mv); self.low_cnt -= 1;
        Self::map_inc(&mut self.mid, mv); self.mid_cnt += 1; self.mid_sum += mv as i64;
      }
    }
  }

  pub fn new(m: i32, k: i32) -> Self {
    MKAverage {
      m: m as usize, k: k as usize,
      window: VecDeque::new(),
      low: BTreeMap::new(), mid: BTreeMap::new(), high: BTreeMap::new(),
      low_cnt: 0, mid_cnt: 0, high_cnt: 0, mid_sum: 0,
    }
  }

  pub fn add_element(&mut self, num: i32) {
    if self.window.len() == self.m {
      let old = self.window.pop_front().unwrap();
      self.remove_rebalance(old);
    }
    self.window.push_back(num);
    self.insert_cascade(num);
  }

  pub fn calculate_mk_average(&self) -> i32 {
    if self.window.len() < self.m { return -1; }
    let denom = (self.m - 2 * self.k) as i64;
    (self.mid_sum / denom) as i32
  }
}