Skip to content

Commit

Permalink
Speed up median
Browse files Browse the repository at this point in the history
Drive-by removal of unwrap in use of `Arc` in `NumbatList::head`
  • Loading branch information
rben01 committed Dec 12, 2024
1 parent b53815d commit a6f2b92
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 8 deletions.
5 changes: 3 additions & 2 deletions numbat/modules/math/statistics.nbt
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,8 @@ fn stdev<D: Dim>(xs: List<D>) -> D = sqrt(variance(xs))
@example("median([1 m, 2 m, 400 cm])")
fn median<D: Dim>(xs: List<D>) -> D = # TODO: this is extremely inefficient
if mod(n, 2) == 1
then element_at((n - 1) / 2, sort(xs))
else mean([element_at(n / 2 - 1, sort(xs)), element_at(n / 2, sort(xs))])
then element_at((n - 1) / 2, sorted)
else (element_at(n / 2 - 1, sorted) + element_at(n / 2, sorted)) / 2
where
n = len(xs)
and sorted = sort(xs)
9 changes: 3 additions & 6 deletions numbat/src/list.rs
Original file line number Diff line number Diff line change
Expand Up @@ -113,12 +113,9 @@ impl<T: Clone> NumbatList<T> {
/// clone the value that's being returned.
pub fn head(self) -> Option<T> {
let front = self.view.map_or(0, |(start, _end)| start);
if Arc::strong_count(&self.alloc) == 1 {
// safety: unwrap cannot fail because we ensured there was only one strong reference above
let mut alloc = Arc::try_unwrap(self.alloc).map_err(|_| ()).unwrap();
alloc.swap_remove_front(front)
} else {
self.alloc.get(front).cloned()
match Arc::try_unwrap(self.alloc) {
Ok(mut solely_owned) => solely_owned.swap_remove_front(front),
Err(shared) => shared.get(front).cloned(),
}
}

Expand Down

0 comments on commit a6f2b92

Please sign in to comment.