Merge pull request #1158 from lsh/dev

[Sparse] Add `triplet_iter_mut()`
This commit is contained in:
Andreas Borgen Longva 2022-10-11 09:00:38 +02:00 committed by GitHub
commit eb2b23b103
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 44 additions and 0 deletions

View File

@ -170,6 +170,16 @@ impl<T> CooMatrix<T> {
.map(|((i, j), v)| (*i, *j, v))
}
/// A mutable iterator over triplets (i, j, v).
// TODO: Consider giving the iterator a concrete type instead of impl trait...?
pub fn triplet_iter_mut(&mut self) -> impl Iterator<Item = (usize, usize, &mut T)> {
self.row_indices
.iter()
.zip(&self.col_indices)
.zip(self.values.iter_mut())
.map(|((i, j), v)| (*i, *j, v))
}
/// Reserves capacity for COO matrix by at least `additional` elements.
///
/// This increase the capacities of triplet holding arrays by reserving more space to avoid

View File

@ -87,6 +87,40 @@ fn coo_construction_for_valid_data() {
}
}
#[test]
fn coo_triplets_iter_mut() {
// Arbitrary matrix, with duplicates
let i = vec![0, 1, 0, 0, 0, 0, 2, 1];
let j = vec![0, 2, 0, 1, 0, 3, 3, 2];
let v = vec![2, 3, 4, 7, 1, 3, 1, 5];
let mut coo =
CooMatrix::<i32>::try_from_triplets(3, 5, i.clone(), j.clone(), v.clone()).unwrap();
let actual_triplets: Vec<_> = coo.triplet_iter_mut().map(|(i, j, v)| (i, j, *v)).collect();
let expected_triplets: Vec<_> = i
.iter()
.zip(&j)
.zip(&v)
.map(|((i, j), v)| (*i, *j, *v))
.collect();
assert_eq!(expected_triplets, actual_triplets);
for (_i, _j, v) in coo.triplet_iter_mut() {
*v += *v;
}
let actual_triplets: Vec<_> = coo.triplet_iter_mut().map(|(i, j, v)| (i, j, *v)).collect();
let v = vec![4, 6, 8, 14, 2, 6, 2, 10];
let expected_triplets: Vec<_> = i
.iter()
.zip(&j)
.zip(&v)
.map(|((i, j), v)| (*i, *j, *v))
.collect();
assert_eq!(expected_triplets, actual_triplets);
}
#[test]
fn coo_try_from_triplets_reports_out_of_bounds_indices() {
{