reduce.h

Go to the documentation of this file.

#pragma once
 
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <span>
#include <utility>
#include <vector>
 
#include "clustering/ndarray.h"
 
namespace clustering::math {

template <class T, Layout L> inline T sum(const NDArray<T, 1, L> &x) noexcept {
  static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
                "sum<T> requires T to be float or double");
  const std::size_t n = x.dim(0);
  T s = T{0};
  for (std::size_t i = 0; i < n; ++i) {
    s += x(i);
  }
  return s;
}

template <class T, Layout L> inline std::size_t argmin(const NDArray<T, 1, L> &x) noexcept {
  static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
                "argmin<T> requires T to be float or double");
  const std::size_t n = x.dim(0);
  assert(n > 0 && "argmin requires non-empty input");
  std::size_t bestIdx = 0;
  T bestVal = x(0);
  for (std::size_t i = 1; i < n; ++i) {
    const T v = x(i);
    if (v < bestVal) {
      bestVal = v;
      bestIdx = i;
    }
  }
  return bestIdx;
}

template <class T, Layout L> inline std::size_t argmax(const NDArray<T, 1, L> &x) noexcept {
  static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
                "argmax<T> requires T to be float or double");
  const std::size_t n = x.dim(0);
  assert(n > 0 && "argmax requires non-empty input");
  std::size_t bestIdx = 0;
  T bestVal = x(0);
  for (std::size_t i = 1; i < n; ++i) {
    const T v = x(i);
    if (v > bestVal) {
      bestVal = v;
      bestIdx = i;
    }
  }
  return bestIdx;
}

template <class T, Layout L>
inline void topk(const NDArray<T, 1, L> &x, std::size_t k, std::span<std::size_t> outIdx) noexcept {
  static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
                "topk<T> requires T to be float or double");
  assert(outIdx.size() == k && "topk requires outIdx.size() == k");
  if (k == 0) {
    return;
  }
  const std::size_t n = x.dim(0);
  assert(k <= n && "topk requires k <= x.dim(0)");
 
  std::vector<std::pair<T, std::size_t>> staged;
  staged.reserve(n);
  for (std::size_t i = 0; i < n; ++i) {
    staged.emplace_back(x(i), i);
  }
 
  std::vector<std::pair<T, std::size_t>> top(k);
  const auto cmp = [](const std::pair<T, std::size_t> &a,
                      const std::pair<T, std::size_t> &b) noexcept {
    // Primary: larger value first. Secondary on ties: smaller index first so the output is
    // deterministic when duplicate values collide.
    if (a.first != b.first) {
      return a.first > b.first;
    }
    return a.second < b.second;
  };
  std::partial_sort_copy(staged.begin(), staged.end(), top.begin(), top.end(), cmp);
 
  for (std::size_t i = 0; i < k; ++i) {
    outIdx[i] = top[i].second;
  }
}
 
} // namespace clustering::math