minimize-malware-spread.rs

struct DSU {
  s: Vec<i32>,
}

impl DSU {
  pub fn new(n: usize) -> Self {
    DSU { s: vec![-1; n] }
  }

  pub fn find(&mut self, i: usize) -> usize {
    if self.s[i] < 0 {
      return i;
    }
    self.s[i] = self.find(self.s[i] as usize) as i32;
    self.s[i] as usize
  }

  pub fn join(&mut self, mut a: usize, mut b: usize) -> bool {
    a = self.find(a);
    b = self.find(b);

    if a == b {
      return false;
    }

    if self.s[a] == self.s[b] {
      self.s[a] -= 1;
    }

    let (par, kid) = match self.s[a] <= self.s[b] {
      true => (a, b),
      false => (b, a),
    };

    self.s[kid] = par as i32;
    self.s[par] -= 1;

    return true;
  }
}

impl Solution {
  pub fn min_malware_spread(graph: Vec<Vec<i32>>, initial: Vec<i32>) -> i32 {
    let mut sets = DSU::new(graph.len());
    graph.iter().enumerate().for_each(|(i, row)| {
      row.iter().enumerate().for_each(|(j, cell)| {
        if *cell == 1 {
          sets.join(i, j);
        }
      })
    });

    let initial_component_freq = initial.iter().fold(vec![0; graph.len()], |mut acc, e| {
      acc[sets.find(*e as usize)] += 1;
      acc
    });
    let compoment_totals = (0..graph.len()).fold(vec![0; graph.len()], |mut acc, e| {
      acc[sets.find(e)] += 1;
      acc
    });

    let saved_parents = (0..graph.len()).map(|i| sets.find(i)).collect::<Vec<usize>>();

    let update_best = |(idx, val): (i32, i32), e: &i32| -> (i32, i32) {
      let tot = compoment_totals[saved_parents[*e as usize]].clone();
      if tot > val || tot == val && *e < idx {
        return (*e, tot);
      }

      (idx, val)
    };

    let (ans, _) = initial
      .iter()
      .filter(|e| initial_component_freq[sets.find(**e as usize)] <= 1)
      .fold((-1, -1), update_best);

    if ans == -1 {
      *initial.iter().min().unwrap()
    } else {
      ans
    }
  }
}