use std::collections::HashMap;
use sprs::{CsMat, TriMat};
#[derive(Debug, Clone)]
pub struct CooBuilder {
rows: usize,
cols: usize,
entries: HashMap<(usize, usize), f64>,
}
impl CooBuilder {
pub fn new(n: usize) -> Self {
Self::new_rect(n, n)
}
pub fn with_capacity(n: usize, capacity: usize) -> Self {
Self::with_capacity_rect(n, n, capacity)
}
pub fn new_rect(rows: usize, cols: usize) -> Self {
Self {
rows,
cols,
entries: HashMap::new(),
}
}
pub fn with_capacity_rect(rows: usize, cols: usize, capacity: usize) -> Self {
Self {
rows,
cols,
entries: HashMap::with_capacity(capacity),
}
}
#[inline]
pub fn n(&self) -> usize {
self.rows
}
#[inline]
pub fn shape(&self) -> (usize, usize) {
(self.rows, self.cols)
}
#[inline]
pub fn add(&mut self, i: usize, j: usize, v: f64) {
if v == 0.0 {
return;
}
debug_assert!(i < self.rows && j < self.cols);
*self.entries.entry((i, j)).or_insert(0.0) += v;
}
#[inline]
pub fn add_sym(&mut self, i: usize, j: usize, v: f64) {
if i == j {
self.add(i, j, v);
} else {
self.add(i, j, v);
self.add(j, i, v);
}
}
pub fn finish_csr(self) -> CsMat<f64> {
let mut tri = TriMat::with_capacity((self.rows, self.cols), self.entries.len());
for ((i, j), v) in self.entries {
if v != 0.0 {
tri.add_triplet(i, j, v);
}
}
tri.to_csr()
}
pub fn finish_csc(self) -> CsMat<f64> {
self.finish_csr().to_csc()
}
}