diff --git a/nalgebra-sparse/src/coo.rs b/nalgebra-sparse/src/coo.rs index 2b302e37..e047745a 100644 --- a/nalgebra-sparse/src/coo.rs +++ b/nalgebra-sparse/src/coo.rs @@ -211,6 +211,26 @@ impl CooMatrix { self.values.push(v); } + /// Remove a single triplet from the matrix. + /// + /// This removes the value `v` from the `i`th row and `j`th column in the matrix. + pub fn clear_triplet(&mut self, i: usize, j: usize, v: T) -> Option<(usize, usize, T)> + where + T: PartialEq, + { + let triple_idx = self + .triplet_iter() + .position(|triplet| triplet == (i, j, &v)); + if let Some(triple_idx) = triple_idx { + self.row_indices.remove(triple_idx); + self.col_indices.remove(triple_idx); + let retv = self.values.remove(triple_idx); + Some((i, j, retv)) + } else { + None + } + } + /// The number of rows in the matrix. #[inline] #[must_use] diff --git a/nalgebra-sparse/tests/unit_tests/coo.rs b/nalgebra-sparse/tests/unit_tests/coo.rs index c70c5f97..6ec88fb1 100644 --- a/nalgebra-sparse/tests/unit_tests/coo.rs +++ b/nalgebra-sparse/tests/unit_tests/coo.rs @@ -226,6 +226,35 @@ fn coo_push_valid_entries() { ); } +#[test] +fn coo_clear_triplet_valid_entries() { + let mut coo = CooMatrix::new(3, 3); + + coo.push(0, 0, 1); + coo.push(0, 0, 2); + coo.push(2, 2, 3); + + // clear a triplet that is not included + let triplet = coo.clear_triplet(0, 0, 0); + assert_eq!(triplet, None); + assert_eq!( + coo.triplet_iter().collect::>(), + vec![(0, 0, &1), (0, 0, &2), (2, 2, &3)] + ); + let triplet = coo.clear_triplet(0, 0, 1); + assert_eq!(triplet, Some((0, 0, 1))); + assert_eq!( + coo.triplet_iter().collect::>(), + vec![(0, 0, &2), (2, 2, &3)] + ); + let triplet = coo.clear_triplet(0, 0, 2); + assert_eq!(triplet, Some((0, 0, 2))); + assert_eq!(coo.triplet_iter().collect::>(), vec![(2, 2, &3)]); + let triplet = coo.clear_triplet(2, 2, 3); + assert_eq!(triplet, Some((2, 2, 3))); + assert_eq!(coo.triplet_iter().collect::>(), vec![]); +} + #[test] fn coo_push_out_of_bounds_entries() { {