Repository: motiwari/BanditPAM (language: python)

Issue:
cosine similarity vs cosine distance
Wanted to say thanks for the repo - this is a giant leap forward in scaling k-medoids.

I noticed that one of the distance metrics available is cosine similarity - not [cosine distance](https://en.wikipedia.org/wiki/Cosine_similarity) which is 1 - cosine similarity.  My intuition tells me  dist(me, me) should be zero, and not one.  

If using cosine similarity is intentional,  I could put in a PR for cos_dist.
minor fix



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/build_linux_wheels.yml b/.github/workflows/build_linux_wheels.yml
index 5d3c43c1..f018d632 100644
--- a/.github/workflows/build_linux_wheels.yml
+++ b/.github/workflows/build_linux_wheels.yml
@@ -20,7 +20,7 @@ jobs:
       - name: Install Ubuntu dependencies
         run: |
           sudo apt update
-          sudo apt install -y build-essential checkinstall libreadline-gplv2-dev libncursesw5-dev libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev 
+          sudo apt install -y build-essential checkinstall libncursesw5-dev libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev 
           sudo apt install -y clang-format cppcheck
 
       # See note above
@@ -41,19 +41,15 @@ jobs:
 
   build_wheels:
     needs: [build_sdist]
-    name: Build wheels on ${{ matrix.os }}
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        os: [ubuntu-20.04]
-
+    name: Build wheels on ubuntu-latest
+    runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
         with:
           submodules: true
 
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.3.1
+        uses: pypa/cibuildwheel@v2.12.0
 
       - uses: actions/upload-artifact@v2
         with:
@@ -76,5 +72,5 @@ jobs:
           # For PyPI
           password: ${{ secrets.PYPI_APIKEY }}
           # For TestPyPI. TODO: RBF
-          # password: ${{ secrets.TEST_PYPI_APIKEY }}
-          # repository_url: https://test.pypi.org/legacy
+#          password: ${{ secrets.TEST_PYPI_APIKEY }}
+#          repository_url: https://test.pypi.org/legacy/
diff --git a/.github/workflows/build_mac_wheels.yml b/.github/workflows/build_mac_wheels.yml
index cea6c888..3bcf049b 100644
--- a/.github/workflows/build_mac_wheels.yml
+++ b/.github/workflows/build_mac_wheels.yml
@@ -20,7 +20,7 @@ jobs:
 
       # Note: CIBW only supports CPython 3.8 and newer for universal2 and arm64 wheels
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.3.1
+        uses: pypa/cibuildwheel@v2.12.0
         env:
           CIBW_ARCHS_MACOS: x86_64 arm64 universal2
           CIBW_BUILD: cp*
@@ -46,5 +46,5 @@ jobs:
           # For PyPI
           password: ${{ secrets.PYPI_APIKEY }}
           # For TestPyPI. TODO: RBF
-          # password: ${{ secrets.TEST_PYPI_APIKEY }}
-          # repository_url: https://test.pypi.org/legacy/
\ No newline at end of file
+#          password: ${{ secrets.TEST_PYPI_APIKEY }}
+#          repository_url: https://test.pypi.org/legacy/
\ No newline at end of file
diff --git a/.github/workflows/run_cmake_build.yml b/.github/workflows/run_cmake_build.yml
new file mode 100644
index 00000000..7c216c05
--- /dev/null
+++ b/.github/workflows/run_cmake_build.yml
@@ -0,0 +1,61 @@
+name: Run CMake build on MacOS
+on: [push, pull_request]
+jobs:
+  build:
+    runs-on: macos-latest
+    strategy:
+      fail-fast: false
+    steps:
+      - uses: actions/checkout@v2
+
+      # On the MacOS machine, when installing from source,
+      # we use LLVM Clang
+      - name: Install MacOS dependencies
+        run: |
+          brew install llvm libomp armadillo
+
+      - name: Install Armadillo 10.7.5+
+        run: |
+          cd ~
+          git clone https://gitlab.com/conradsnicta/armadillo-code.git
+          cd armadillo-code
+          sudo cmake .
+          sudo make install
+
+      - name: Install carma
+        run: |
+          cd ~
+          git clone https://github.com/RUrlus/carma.git
+          cd carma
+          mkdir build
+          cd build
+          sudo cmake -DCARMA_INSTALL_LIB=ON ..
+          sudo cmake --build . --config Release --target install
+          cd ~
+
+      - name: Downloading data files for tests
+        run: |
+          mkdir -p data
+          curl -XGET https://motiwari.com/banditpam_data/MNIST_1k.csv > data/MNIST_1k.csv
+
+      # On the MacOS machine, when installing from source,
+      # we use LLVM Clang
+      # TODO(@motiwari): Clean this environment variables up
+      - name: Verify that the C++ executable compiles and runs
+        run : |
+          echo 'export PATH="/usr/local/opt/llvm/bin:$PATH"' >> /Users/runner/.bash_profile
+          source /Users/runner/.bash_profile
+          export LDFLAGS="-L/usr/local/opt/llvm/lib, -L/usr/local/opt/llvm/lib/c++ -Wl,-rpath,/usr/local/opt/llvm/lib/c++"
+          export CPPFLAGS="-I/usr/local/opt/llvm/include"
+          export DYLD_LIBRARY_PATH=/usr/local/opt/libomp/lib:/usr/local/opt/libomp/include
+          export CC=/usr/local/opt/llvm/bin/clang
+          export CXX=/usr/local/opt/llvm/bin/clang++
+          git clone https://github.com/ThrunGroup/BanditPAM
+          cd BanditPAM
+          mkdir build
+          cd build
+          sudo CC=/usr/local/opt/llvm/bin/clang CXX=/usr/local/opt/llvm/bin/clang++ cmake ..
+          sudo make
+          src/BanditPAM -f ../data/MNIST_1k.csv -k 5 -l L2
+        env:
+          CPLUS_INCLUDE_PATH: /usr/local/include/carma:/usr/local/Cellar/libomp/15.0.2/include:/usr/local/Cellar/libomp/15.0.7/include
\ No newline at end of file
diff --git a/.github/workflows/run_style_checks.yml b/.github/workflows/run_style_checks.yml
index dde6e354..9afe1e26 100644
--- a/.github/workflows/run_style_checks.yml
+++ b/.github/workflows/run_style_checks.yml
@@ -6,28 +6,32 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.6", "3.7", "3.8", "3.9"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
+
       - name: Install style check dependencies
         run: |
           sudo apt update
           sudo apt install -y clang-format cppcheck
-      - name: Install Python dependencies 
+
+      - name: Install Python dependencies
         run: |
           python -m pip install --upgrade pip
           pip install flake8 pytest cpplint
           pip install -r requirements.txt
+
       - name: Lint with flake8
         run: |
           # stop the build if there are Python syntax errors or undefined names
           flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
           # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
           flake8 . --count --exit-zero --max-complexity=10 --max-line-length=79 --statistics --exclude headers/carma
+
       - name: Lint with cpplint
         run: |
           cpplint --exclude=headers/carma --exclude=build --filter=-legal/copyright,-build/c++11,-build/include_subdir --recursive .
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 92fed4b4..d76e0d0c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -28,3 +28,4 @@ tmp/
 var/
 dist/
 wheelhouse/
+.eggs/
diff --git a/CMakeLists.txt b/CMakeLists.txt
index a5fb802f..52e49a74 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -2,14 +2,18 @@ cmake_minimum_required(VERSION 3.17)
 
 project(BanditPAM VERSION 1.0 LANGUAGES CXX)
 
-set(CMAKE_BUILD_TYPE Release)
+# TODO(@motiwari): Should RELEASE? be in caps
+set(CMAKE_BUILD_TYPE Debug)
 set(CMAKE_CXX_STANDARD 17)
 
+set(CMAKE_CXX_FLAGS_RELEASE "-O3")
+
 # Note: ThreadSanitizer and AddressSanitizer can give spurious errors around the #pragma omp critical.
 # They don't really matter because the threads always resolve with the same end state
 # ThreadSanitizer also introduces a very large overhead, AddressSanitizer is more reasonable
 set(CMAKE_CXX_FLAGS_DEBUG "-Wall -Wextra -g -fno-omit-frame-pointer -fsanitize=address")
 
-set(CMAKE_CXX_FLAGS_RELEASE "-O3")
+# NOTE: Need to explicitly pass -O0 to enable printing of armadillo matrices
+#set(CMAKE_CXX_FLAGS_DEBUG "-O0 -Wall -Wextra -g -fno-omit-frame-pointer")
 
 add_subdirectory(src)
diff --git a/docs/headerFile b/docs/headerFile
index 69230c41..f7af07cb 100644
--- a/docs/headerFile
+++ b/docs/headerFile
@@ -6,7 +6,6 @@
 <meta http-equiv="X-UA-Compatible" content="IE=11"/>
 <meta name="generator" content="Doxygen $doxygenversion"/>
 <meta name="viewport" content="width=device-width, initial-scale=1"/>
-<link rel="shortcut icon" href="$relpath^banditpam_favicon.ico" type="image/x-icon" />
 <!--BEGIN PROJECT_NAME--><title>$projectname: $title</title><!--END PROJECT_NAME-->
 <!--BEGIN !PROJECT_NAME--><title>$title</title><!--END !PROJECT_NAME-->
 <link href="$relpath^tabs.css" rel="stylesheet" type="text/css"/>
diff --git a/docs/sphinx/conf.py b/docs/sphinx/conf.py
index f931aca8..c978201f 100644
--- a/docs/sphinx/conf.py
+++ b/docs/sphinx/conf.py
@@ -20,12 +20,12 @@
 
 # -- Project information -----------------------------------------------------
 
-project = 'BanditPAM'
-copyright = '2021, Mo Tiwari'
-author = 'Mo Tiwari'
+project = "BanditPAM"
+copyright = "2021, Mo Tiwari"
+author = "Mo Tiwari"
 
 # The full version, including alpha/beta/rc tags
-release = '2.0.0'
+release = "2.0.0"
 
 
 # -- General configuration ---------------------------------------------------
@@ -34,16 +34,16 @@
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
-    'myst_parser',
+    "myst_parser",
 ]
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
+exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
 
 
 # -- Options for HTML output -------------------------------------------------
@@ -51,20 +51,20 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = 'alabaster'
+html_theme = "alabaster"
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+html_static_path = ["_static"]
 
 # -- Bridge with Doxygen -------------------------------------------------
 
 # TODO(@motiwari): Unsafe subprocess calls, hacky copy-over of index.html
-subprocess.call('cd .. ; doxygen Doxyfile', shell=True)
+subprocess.call("cd .. ; doxygen Doxyfile", shell=True)
 readme_file = glob.glob("../html/*_README.html")[0]
-subprocess.call('cp {} ../html/index.html'.format(readme_file), shell=True)
+subprocess.call("cp {} ../html/index.html".format(readme_file), shell=True)
 
-html_extra_path = ['../html']
+html_extra_path = ["../html"]
 
-html_favicon = 'banditpam_favicon.ico'
+html_favicon = "banditpam_favicon.ico"
diff --git a/headers/algorithms/banditpam.hpp b/headers/algorithms/banditpam.hpp
index eb518749..6489de89 100644
--- a/headers/algorithms/banditpam.hpp
+++ b/headers/algorithms/banditpam.hpp
@@ -18,10 +18,12 @@ class BanditPAM : public km::KMedoids {
  public:
   /**
    * @brief Runs BanditPAM to identify a dataset's medoids.
-   * 
+   *
    * @param inputData Input data to cluster
    */
-  void fitBanditPAM(const arma::fmat& inputData);
+  void fitBanditPAM(
+    const arma::fmat& inputData,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat);
 
   /**
    * @brief Empirical estimation of standard deviation of arm returns
@@ -31,28 +33,30 @@ class BanditPAM : public km::KMedoids {
    * @param bestDistances Contains best distances from each point to medoids
    * @param useAbsolute Flag to use the absolute distance to each arm instead
    * of improvement over prior loss; necessary for the first BUILD step
-   * 
+   *
    * @returns Estimate of each arm's standard deviation
    */
   arma::frowvec buildSigma(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     const arma::frowvec& bestDistances,
     const bool useAbsolute);
 
-  /** 
+  /**
    * @brief Estimates the mean reward for each arm in the BUILD steps.
-   * 
+   *
    * @param data Transposed input data to cluster
    * @param target Candidate datapoints to consider adding as medoids
    * @param bestDistances Contains best distances from each point to medoids
    * @param useAbsolute Flag to use the absolute distance to each arm instead
    * of improvement over prior loss; necessary for the first BUILD step
    * @param exact 0 if using standard batch size; size of dataset otherwise
-   * 
+   *
    * @returns Estimate of each arm's change in loss
    */
   arma::frowvec buildTarget(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     const arma::uvec* target,
     const arma::frowvec* bestDistances,
     const bool useAbsolute,
@@ -63,7 +67,7 @@ class BanditPAM : public km::KMedoids {
    *
    * Draws batch size reference points with replacement and uses the estimated
    * reward of adding candidate medoids to the set of medoids. Constructs
-   * confidence intervals for expected loss when adding each candidate as a 
+   * confidence intervals for expected loss when adding each candidate as a
    * medoid, and continues drawing reference points until the best candidate's
    * confidence interval is disjoint from all others.
    *
@@ -74,6 +78,7 @@ class BanditPAM : public km::KMedoids {
    */
   void build(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     arma::urowvec* medoidIndices,
     arma::fmat* medoids);
 
@@ -85,11 +90,12 @@ class BanditPAM : public km::KMedoids {
    * @param bestDistances Contains best distances from each point to medoids
    * @param secondBestDistances Contains second best distances from each point to medoids
    * @param assignments Assignments of datapoints to their closest medoid
-   * 
+   *
    * @returns Estimate of each arm's standard deviation
    */
   arma::fmat swapSigma(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     const arma::frowvec* bestDistances,
     const arma::frowvec* secondBestDistances,
     const arma::urowvec* assignments);
@@ -105,11 +111,12 @@ class BanditPAM : public km::KMedoids {
    * @param secondBestDistances Contains second best distances from each point to medoids
    * @param assignments Assignments of datapoints to their closest medoid
    * @param exact 0 if using standard batch size; size of dataset otherwise
-   * 
+   *
    * @returns Estimate of each arm's change in loss
    */
-  arma::fvec swapTarget(
+  arma::fmat swapTarget(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     const arma::urowvec* medoidIndices,
     const arma::uvec* targets,
     const arma::frowvec* bestDistances,
@@ -119,10 +126,10 @@ class BanditPAM : public km::KMedoids {
 
   /**
   * @brief Performs the SWAP step of BanditPAM.
-  * 
+  *
   * Draws batch size reference points with replacement and uses the estimated
   * reward of performing a (medoid, non-medoid) swap. Constructs
-  * confidence intervals for expected loss when performing the swap, 
+  * confidence intervals for expected loss when performing the swap,
   * and continues drawing reference points until the best candidate's
   * confidence interval is disjoint from all others.
   *
@@ -135,6 +142,7 @@ class BanditPAM : public km::KMedoids {
   */
   void swap(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     arma::urowvec* medoidIndices,
     arma::fmat* medoids,
     arma::urowvec* assignments);
diff --git a/headers/algorithms/banditpam_orig.hpp b/headers/algorithms/banditpam_orig.hpp
new file mode 100644
index 00000000..5165cdde
--- /dev/null
+++ b/headers/algorithms/banditpam_orig.hpp
@@ -0,0 +1,152 @@
+#ifndef HEADERS_ALGORITHMS_BANDITPAM_ORIG_HPP_
+#define HEADERS_ALGORITHMS_BANDITPAM_ORIG_HPP_
+
+#include <omp.h>
+#include <armadillo>
+#include <vector>
+#include <fstream>
+#include <iostream>
+
+
+#include "kmedoids_algorithm.hpp"
+
+namespace km {
+/**
+ * @brief Contains all necessary BanditPAM_orig functions
+ */
+class BanditPAM_orig : public km::KMedoids {
+ public:
+  /**
+   * @brief Runs BanditPAM_orig to identify a dataset's medoids.
+   *
+   * @param inputData Input data to cluster
+   */
+  void fitBanditPAM_orig(
+      const arma::fmat& inputData,
+      std::optional<std::reference_wrapper<const arma::fmat>> distMat);
+
+  /**
+   * @brief Empirical estimation of standard deviation of arm returns
+   * in the BUILD step.
+   *
+   * @param data Transposed input data to cluster
+   * @param bestDistances Contains best distances from each point to medoids
+   * @param useAbsolute Flag to use the absolute distance to each arm instead
+   * of improvement over prior loss; necessary for the first BUILD step
+   *
+   * @returns Estimate of each arm's standard deviation
+   */
+  arma::frowvec buildSigma(
+    const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+    const arma::frowvec& bestDistances,
+    const bool useAbsolute);
+
+  /**
+   * @brief Estimates the mean reward for each arm in the BUILD steps.
+   *
+   * @param data Transposed input data to cluster
+   * @param target Candidate datapoints to consider adding as medoids
+   * @param bestDistances Contains best distances from each point to medoids
+   * @param useAbsolute Flag to use the absolute distance to each arm instead
+   * of improvement over prior loss; necessary for the first BUILD step
+   * @param exact 0 if using standard batch size; size of dataset otherwise
+   *
+   * @returns Estimate of each arm's change in loss
+   */
+  arma::frowvec buildTarget(
+    const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+    const arma::uvec* target,
+    const arma::frowvec* bestDistances,
+    const bool useAbsolute,
+    const size_t exact);
+
+  /**
+   * @brief Performs the BUILD step of BanditPAM_orig.
+   *
+   * Draws batch size reference points with replacement and uses the estimated
+   * reward of adding candidate medoids to the set of medoids. Constructs
+   * confidence intervals for expected loss when adding each candidate as a
+   * medoid, and continues drawing reference points until the best candidate's
+   * confidence interval is disjoint from all others.
+   *
+   * @param data Transposed input data to cluster
+   * @param medoidIndices Array of medoids that is modified in place
+   * as medoids are identified
+   * @param medoids Matrix that contains the coordinates of each medoid
+   */
+  void build(
+    const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+    arma::urowvec* medoidIndices,
+    arma::fmat* medoids);
+
+  /**
+   * @brief Empirical estimation of standard deviation of arm returns
+   * in the SWAP step.
+   *
+   * @param data Transposed input data to cluster
+   * @param bestDistances Contains best distances from each point to medoids
+   * @param secondBestDistances Contains second best distances from each point to medoids
+   * @param assignments Assignments of datapoints to their closest medoid
+   *
+   * @returns Estimate of each arm's standard deviation
+   */
+  arma::fmat swapSigma(
+    const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+    const arma::frowvec* bestDistances,
+    const arma::frowvec* secondBestDistances,
+    const arma::urowvec* assignments);
+
+  /**
+   * @brief Estimates the mean reward for each arm in SWAP step.
+   *
+   * @param data Transposed input data to cluster
+   * @param medoidIndices Array of medoids that is modified in place
+   * as medoids are swapped in and out
+   * @param targets Set of potential swaps to be evaluated
+   * @param bestDistances Contains best distances from each point to medoids
+   * @param secondBestDistances Contains second best distances
+   * from each point to medoids
+   * @param assignments Assignments of datapoints to their closest medoid
+   * @param exact 0 if using standard batch size; size of dataset otherwise
+   *
+   * @returns Estimate of each arm's change in loss
+   */
+  arma::fvec swapTarget(
+    const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+    const arma::urowvec* medoidIndices,
+    const arma::uvec* targets,
+    const arma::frowvec* bestDistances,
+    const arma::frowvec* secondBestDistances,
+    const arma::urowvec* assignments,
+    const size_t exact);
+
+  /**
+  * @brief Performs the SWAP step of BanditPAM_orig.
+  *
+  * Draws batch size reference points with replacement and uses the estimated
+  * reward of performing a (medoid, non-medoid) swap. Constructs
+  * confidence intervals for expected loss when performing the swap,
+  * and continues drawing reference points until the best candidate's
+  * confidence interval is disjoint from all others.
+  *
+  * @param data Transposed input data to cluster
+  * @param medoidIndices Array of medoid indices created from the BUILD step
+  * that is modified in place as better medoids are identified
+  * @param medoids Matrix of possible medoids that is updated as the bandit
+  * learns which datapoints will be unlikely to be good candidates
+  * @param assignments Array of containing the medoid each point is closest to
+  */
+  void swap(
+    const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+    arma::urowvec* medoidIndices,
+    arma::fmat* medoids,
+    arma::urowvec* assignments);
+};
+}  // namespace km
+#endif  // HEADERS_ALGORITHMS_BANDITPAM_ORIG_HPP_
diff --git a/headers/algorithms/fastpam1.hpp b/headers/algorithms/fastpam1.hpp
index e06d1790..a633196e 100644
--- a/headers/algorithms/fastpam1.hpp
+++ b/headers/algorithms/fastpam1.hpp
@@ -1,11 +1,11 @@
 #ifndef HEADERS_ALGORITHMS_FASTPAM1_HPP_
 #define HEADERS_ALGORITHMS_FASTPAM1_HPP_
 
-#include <omp.h>
 #include <armadillo>
-#include <vector>
-#include <fstream>
+#include <omp.h>
 #include <iostream>
+#include <fstream>
+#include <vector>
 
 #include "kmedoids_algorithm.hpp"
 
@@ -21,7 +21,9 @@ class FastPAM1 : public km::KMedoids {
    *
    * @param inputData Input data to cluster
    */
-  void fitFastPAM1(const arma::fmat& inputData);
+  void fitFastPAM1(
+    const arma::fmat& inputData,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat);
 
   /**
    * @brief Performs the BUILD step of FastPAM1.
@@ -36,6 +38,7 @@ class FastPAM1 : public km::KMedoids {
    */
   void buildFastPAM1(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     arma::urowvec* medoidIndices);
 
   /**
@@ -53,6 +56,7 @@ class FastPAM1 : public km::KMedoids {
    */
   void swapFastPAM1(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     arma::urowvec* medoidIndices,
     arma::urowvec* assignments);
 };
diff --git a/headers/algorithms/kmedoids_algorithm.hpp b/headers/algorithms/kmedoids_algorithm.hpp
index 9e1ec293..83206c66 100644
--- a/headers/algorithms/kmedoids_algorithm.hpp
+++ b/headers/algorithms/kmedoids_algorithm.hpp
@@ -3,6 +3,7 @@
 
 #include <omp.h>
 #include <armadillo>
+#include <optional>
 #include <vector>
 #include <fstream>
 #include <iostream>
@@ -24,12 +25,19 @@ namespace km {
  */
 class KMedoids {
  public:
+  // NOTE: The order of arguments in this constructor must match that of the
+  //  arguments in kmedoids_pywrapper.cpp, otherwise undefined behavior can
+  //  result (variables being initialized with others' values)
   KMedoids(
     size_t nMedoids = 5,
     const std::string& algorithm = "BanditPAM",
-    size_t maxIter = 1000,
-    size_t buildConfidence = 1000,
-    size_t swapConfidence = 10000,
+    size_t maxIter = 100,
+    size_t buildConfidence = 3,
+    size_t swapConfidence = 4,
+    bool useCache = true,
+    bool usePerm = true,
+    size_t cacheWidth = 1000,
+    bool parallelize = true,
     size_t seed = 0);
 
   ~KMedoids();
@@ -42,7 +50,10 @@ class KMedoids {
    * 
    * @throws if the input data is empty.
    */
-  void fit(const arma::fmat& inputData, const std::string& loss);
+  void fit(
+    const arma::fmat& inputData,
+    const std::string& loss,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat);
 
   /**
    * @brief Returns the medoids at the end of the BUILD step.
@@ -190,8 +201,124 @@ class KMedoids {
    */
   float getAverageLoss() const;
 
-  /// The cache will be of size cacheMultiplier*nlogn
-  size_t cacheMultiplier = 1000;
+  /**
+   * @brief Returns whether a distance cache is being used
+   *
+   * @return Whether a distance cache is being used
+   */
+  bool getUseCache() const;
+
+  /**
+   * @brief Sets whether a distance cache should be used
+   *
+   * @param newUseCache Whether to use a distance cache
+   */
+  void setUseCache(bool newUseCache);
+
+  /**
+   * @brief Returns whether a permutation of reference points is being used
+   *
+   * @return Whether a permutation of reference points is being used
+   */
+  bool getUsePerm() const;
+
+  /**
+   * @brief Sets whether a permutation of reference points should used
+   *
+   * @param newUsePerm Whether a permutation of reference points should used
+   */
+  void setUsePerm(bool newUsePerm);
+
+  /**
+   * @brief Returns the cache width being used
+   *
+   * @return The cache width being used
+   */
+  size_t getCacheWidth() const;
+
+  /**
+   * @brief Sets the new cache width to use
+   *
+   * @param newCacheWidth The new cache width to use
+   */
+  void setCacheWidth(size_t newCacheWidth);
+
+  /**
+   * @brief Whether the algorithm is parallelized via OpenMP
+   *
+   * @return Whether the algorithm is being parallelized via OpenMP
+   */
+  bool getParallelize() const;
+
+  /**
+   * @brief Whether to parallelize the algorithm via OpenMP
+   *
+   * @param newParallelize Whether to parallelize the algorithm via OpenMP
+   */
+  void setParallelize(bool newParallelize);
+
+  /**
+   * @brief Get total sample complexity of .fit() call
+   *
+   * @return Total sample complexity of last .fit() call
+   */
+  size_t getDistanceComputations(const bool includeMisc = false) const;
+
+  /**
+   * @brief Get number of miscellaneous distance computations
+   *
+   * @return Total number of miscellaneous distance computations
+   */
+  size_t getMiscDistanceComputations() const;
+
+  /**
+   * @brief Get total sample complexity of BUILD step
+   *
+   * @return Total sample complexity of last BUILD step
+   */
+  size_t getBuildDistanceComputations() const;
+
+  /**
+   * @brief Get total sample complexity of SWAP steps
+   *
+   * @return Total sample complexity of last SWAP steps
+   */
+  size_t getSwapDistanceComputations() const;
+
+  /**
+   * @brief Get number of times we wrote to the cache
+   *
+   * @return Number of times we wrote to the cache
+   */
+  size_t getCacheWrites() const;
+
+  /**
+   * @brief Get number of cache hits
+   *
+   * @return Number of cache hits
+   */
+  size_t getCacheHits() const;
+
+  /**
+   * @brief Get number of cache misses
+   *
+   * @return Number of cache misses
+   */
+  size_t getCacheMisses() const;
+
+  /**
+   * @brief Get total number of milliseconds for the whole SWAP procedure
+   *
+   * @return Total number of milliseconds for the whole SWAP procedure
+   */
+  size_t getTotalSwapTime() const;
+
+  /**
+   * @brief Get average number of milliseconds per swap step
+   *
+   * @return Average number of milliseconds per swap step
+   */
+  float getTimePerSwap() const;
 
   /// The cache which stores pairwise distance computations
   float* cache;
@@ -205,11 +332,9 @@ class KMedoids {
   /// A map from permutation index of each point to its original index
   std::unordered_map<size_t, size_t> reindex;
 
-  /// Used for debugging only to toggle a fixed permutation of points
-  bool usePerm = true;
+  /// Determines whether we use a user-provided distance matrix
+  bool useDistMat = false;
 
-  /// Used for debugging only to toggle use of the cache
-  bool useCacheP = true;
 
  protected:
   /**
@@ -226,6 +351,7 @@ class KMedoids {
    */
   void calcBestDistancesSwap(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     const arma::urowvec* medoidIndices,
     arma::frowvec* bestDistances,
     arma::frowvec* secondBestDistances,
@@ -243,26 +369,29 @@ class KMedoids {
    */
   float calcLoss(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     const arma::urowvec* medoidIndices);
 
   /**
    * @brief A wrapper around the given loss function that caches distances
    * between the given points.
    * 
-   * NOTE: if you change useCache, also change useCacheP
+   * NOTE: if you change useCacheFunctionOverride, also change useCache
    * 
    * @param data Transposed data to cluster
    * @param i Index of first datapoint
    * @param j Index of second datapoint
-   * @param useCache Indices of the medoids in the dataset
+   * @param useCacheFunctionOverride Whether to use the cache in this function (by default, uses value of useCache)
    * 
    * @returns The distance between points i and j
    */
   float cachedLoss(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     const size_t i,
     const size_t j,
-    const bool useCache = true);
+    const size_t category,
+    const bool useCacheFunctionOverride = true);
 
   /// If using an L_p loss, the value of p
   size_t lp;
@@ -363,25 +492,64 @@ class KMedoids {
     const;
 
   /// Number of SWAP steps performed
-  size_t steps;
+  size_t steps = 0;
 
   /// Governs the error rate of each BUILD step in BanditPAM
-  size_t buildConfidence;
+  size_t buildConfidence = 1;
 
   /// Governs the error rate of each SWAP step in BanditPAM
-  size_t swapConfidence;
+  size_t swapConfidence = 1;
+
+  /// Used for debugging only to toggle use of the cache
+  bool useCache = true;
+
+  /// Used for debugging only to toggle a fixed permutation of points
+  bool usePerm = true;
+
+  /// The cache will be of size cacheWidth*n
+  size_t cacheWidth = 1000;
+
+  /// Determines whether we parallelize the algorithm with OpenMP
+  bool parallelize = true;
+
+  /// The random seed with which to perform the clustering
+  size_t seed = 0;
 
   /// Used for floatcomparisons, primarily number of "arms" remaining
   const float precision = 0.001;
 
   /// Contains the average loss at the last step of the algorithm
-  float averageLoss;
+  float averageLoss = 0.0;
 
   /// Number of points to sample per reference batch
   size_t batchSize = 100;
 
-  /// The random seed with which to perform the clustering
-  size_t seed = 0;
+  /// The number of non-cache distance computations we compute
+  /// in the BUILD step. For debugging only.
+  size_t numMiscDistanceComputations = 0;
+
+  /// The number of non-cache distance computations we compute
+  /// in the BUILD step. For debugging only.
+  size_t numBuildDistanceComputations = 0;
+
+  /// The number of non-cache distance computations we compute.
+  /// For debugging only.
+  size_t numSwapDistanceComputations = 0;
+
+  /// The number of cache hits (distance computations we reuse).
+  /// For debugging only.
+  size_t numCacheWrites = 0;
+
+  /// The number of cache writes (distance computations we save).
+  /// For debugging only.
+  size_t numCacheHits = 0;
+
+  /// The number of cache misses, i.e., distance computations we
+  /// need to compute. For debugging only.
+  size_t numCacheMisses = 0;
+
+  /// The number of milliseconds taken per swap step, on average
+  size_t totalSwapTime = 0;
 };
 }  // namespace km
 #endif  // HEADERS_ALGORITHMS_KMEDOIDS_ALGORITHM_HPP_
diff --git a/headers/algorithms/pam.hpp b/headers/algorithms/pam.hpp
index 905b0d95..1b5340e4 100644
--- a/headers/algorithms/pam.hpp
+++ b/headers/algorithms/pam.hpp
@@ -10,7 +10,6 @@
 #include "kmedoids_algorithm.hpp"
 
 
-
 namespace km {
 /**
  * @brief Contains all necessary PAM functions
@@ -18,11 +17,13 @@ namespace km {
 class PAM : public km::KMedoids {
  public:
   /**
-   * @brief Runs PAM to identify a dataset's medoids.
-   * 
-   * @param inputData Input data to cluster
-   */
-  void fitPAM(const arma::fmat& inputData);
+  * @brief Runs PAM to identify a dataset's medoids.
+  *
+  * @param inputData Input data to cluster
+  */
+  void fitPAM(
+  const arma::fmat& inputData,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat);
 
   /**
   * @brief Performs the BUILD step of PAM.
@@ -37,14 +38,15 @@ class PAM : public km::KMedoids {
   */
   void buildPAM(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     arma::urowvec* medoidIndices);
 
-  /** 
+  /**
   * @brief Performs the SWAP steps of BanditPAM.
-  * 
+  *
   * Loops over all (medoid, non-medoid) pairs and computes the change in loss
   * when the points are swapped in and out of the medoid set. Then updates
-  * the list of medoids by performing the swap that would lower the overall 
+  * the list of medoids by performing the swap that would lower the overall
   * loss the most, provided at least one such swap would reduce the loss.
   *
   * @param data Transposed input data to cluster
@@ -54,6 +56,7 @@ class PAM : public km::KMedoids {
   */
   void swapPAM(
     const arma::fmat& data,
+    std::optional<std::reference_wrapper<const arma::fmat>> distMat,
     arma::urowvec* medoidIndices,
     arma::urowvec* assignments);
 };
diff --git a/headers/python_bindings/kmedoids_pywrapper.hpp b/headers/python_bindings/kmedoids_pywrapper.hpp
index 847331a9..ede77534 100644
--- a/headers/python_bindings/kmedoids_pywrapper.hpp
+++ b/headers/python_bindings/kmedoids_pywrapper.hpp
@@ -2,6 +2,7 @@
 #define HEADERS_PYTHON_BINDINGS_KMEDOIDS_PYWRAPPER_HPP_
 
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 #include <pybind11/numpy.h>
 #include <carma>
 #include <armadillo>
@@ -71,6 +72,57 @@ class KMedoidsWrapper : public km::KMedoids {
    * nearest medoid
    */
   float getLossPython();
+
+  /**
+   * @brief Returns the number of distance computations (sample complexity) used by .fit()
+   *
+   * The total number of distance computations (sample complexity) used by .fit()
+   */
+  size_t getDistanceComputationsPython(const bool includeMisc);
+
+  // TODO(@motiwari): Add docstring
+  size_t getMiscDistanceComputationsPython();
+
+  // TODO(@motiwari): Add docstring
+  size_t getBuildDistanceComputationsPython();
+
+  // TODO(@motiwari): Add docstring
+  size_t getSwapDistanceComputationsPython();
+
+  /**
+   * @brief Returns the number of cache writes done by .fit()
+   *
+   * The number of cache writes performed by the last call to .fit()
+   */
+  size_t getCacheWritesPython();
+
+  /**
+   * @brief Returns the number of cache hits from the last call to .fit()
+   *
+   * The number of cache hits from the last call to .fit()
+  */
+  size_t getCacheHitsPython();
+
+  /**
+   * @brief Returns the number of cache misses by the last call to .fit()
+   *
+   * The number of cache misses from the last call to .fit()
+   */
+  size_t getCacheMissesPython();
+
+  /**
+   * @brief Returns the total time for the entire SWAP procedure by the last call to .fit()
+   *
+   * The total time for the entire SWAP procedure by the last call to .fit()
+   */
+  size_t getTotalSwapTimePython();
+
+  /**
+   * @brief Returns the average time per swap step by the last call to .fit()
+   *
+   * The average time per swap step by the last call to .fit()
+   */
+  float getTimePerSwapPython();
 };
 
 // TODO(@motiwari): Encapsulate these
@@ -104,5 +156,45 @@ void steps_python(pybind11::class_<km::KMedoidsWrapper> *);
  * @brief Binding for the C++ function KMedoids::calcLoss()
  */
 void loss_python(pybind11::class_<km::KMedoidsWrapper> *);
+
+/**
+ * @brief Binding for the C++ function KMedoids::getDistanceComputations()
+ */
+void distance_computations_python(pybind11::class_<km::KMedoidsWrapper> *);
+
+// TODO(@motiwari): Add docstring
+void misc_distance_computations_python(pybind11::class_<km::KMedoidsWrapper> *);
+
+// TODO(@motiwari): Add docstring
+void build_distance_computations_python(
+    pybind11::class_<km::KMedoidsWrapper> *);
+
+// TODO(@motiwari): Add docstring
+void swap_distance_computations_python(pybind11::class_<km::KMedoidsWrapper> *);
+
+/**
+ * @brief Binding for the C++ function KMedoids::getCacheWrites()
+ */
+void cache_writes_python(pybind11::class_<km::KMedoidsWrapper> *);
+
+/**
+ * @brief Binding for the C++ function KMedoids::getCacheHits()
+ */
+void cache_hits_python(pybind11::class_<km::KMedoidsWrapper> *);
+
+/**
+ * @brief Binding for the C++ function KMedoids::getCacheMisses()
+ */
+void cache_misses_python(pybind11::class_<km::KMedoidsWrapper> *);
+
+/**
+ * @brief Binding for the C++ function KMedoids::getTotalSwapTime()
+ */
+void total_swap_time_python(pybind11::class_<km::KMedoidsWrapper> *);
+
+/**
+ * @brief Binding for the C++ function KMedoids::getTimePerSwap()
+ */
+void time_per_swap_python(pybind11::class_<km::KMedoidsWrapper> *);
 }  // namespace km
 #endif  // HEADERS_PYTHON_BINDINGS_KMEDOIDS_PYWRAPPER_HPP_
diff --git a/scripts/cache_measurements.py b/scripts/cache_measurements.py
new file mode 100644
index 00000000..b3d29804
--- /dev/null
+++ b/scripts/cache_measurements.py
@@ -0,0 +1,188 @@
+import numpy as np
+import time
+import os
+from banditpam import KMedoids
+import math
+
+from tests.constants import MILLISECONDS_IN_A_SECOND
+
+
+def time_measured_fit(
+    kmed: KMedoids,
+    X: np.array,
+    loss: str = "L2",
+):
+    start = time.time()
+    kmed.fit(X, loss)
+    elapsed_time = time.time() - start  # Elapsed time in seconds
+    return elapsed_time
+
+
+def get_cache_statistics(
+    kmed: KMedoids,
+    X: np.array,
+    loss: str = "L2",
+    cache_width: int = 1000,
+):
+    kmed.cache_width = cache_width
+    time = time_measured_fit(kmed, X, loss)
+    print("Cache Width: {:,}".format(kmed.cache_width))
+    print("Distance Computations: {:,}".format(kmed.distance_computations))
+    print("Swap Steps: {:,}".format(kmed.steps))
+    print("Cache Writes: {:,}".format(kmed.cache_writes))
+    print("Cache Hits: {:,}".format(kmed.cache_hits))
+    print("Cache Misses: {:,}".format(kmed.cache_misses))
+    return (
+        time,
+        kmed.cache_width,
+        kmed.steps,
+        kmed.distance_computations,
+        kmed.cache_writes,
+        kmed.cache_hits,
+        kmed.cache_misses,
+    )
+
+
+def main():
+    def test_cache_stats():
+        X = np.loadtxt(os.path.join("data", "MNIST_1k.csv"))
+        kmed = KMedoids(n_medoids=5, algorithm="BanditPAM")
+
+        (
+            time_1000,
+            width_1000,
+            writes_1000,
+            hits_1000,
+            misses_1000,
+        ) = get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=1000)
+        time_750, width_750, writes_750, hits_750, misses_750 = \
+            get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=750)
+        time_500, width_500, writes_500, hits_500, misses_500 = \
+            get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=500)
+        time_250, width_250, writes_250, hits_250, misses_250 = \
+            get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=250)
+        time_0, width_0, writes_0, hits_0, misses_0 = \
+            get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=0)
+
+        assert (
+            hits_1000 > hits_750 > hits_500 > hits_250 > hits_0
+        ), "Cache hits should increase as cache size increases"
+        assert (
+            misses_1000 < misses_750 < misses_500 < misses_250 < misses_0
+        ), "Cache misses should decrease as cache size increases"
+        assert (
+            misses_1000 == 0
+        ), "There should be no cache misses at width=1000 as the" \
+           " whole dataset fits in memory"
+        assert (
+            hits_0 == 0
+        ), "There should be no cache hits at width=0 as there is no cache"
+        assert (
+            writes_0 == 0
+        ), "There should be no cache writes at width=0 as there is no cache"
+        assert width_0 == 0, "Cache width should be 0 when set to 0"
+        assert width_250 == 250, "Cache width should be 250 when set to 250"
+        assert width_500 == 500, "Cache width should be 500 when set to 500"
+        assert width_750 == 750, "Cache width should be 750 when set to 750"
+        assert width_1000 == 1000, "Cache width should be 1000 when set " \
+                                   "to 1000"
+
+    def test_parallelization():
+        X = np.loadtxt(os.path.join("data", "MNIST_10k.csv"))
+        kmed = KMedoids(n_medoids=5, algorithm="BanditPAM")
+        time_parallel = time_measured_fit(kmed=kmed, X=X, loss="L2")
+        kmed.parallelize = False
+        time_no_parallel = time_measured_fit(kmed=kmed, X=X, loss="L2")
+
+        print(time_no_parallel, time_parallel)
+        assert (
+            time_no_parallel > time_parallel
+        ), "Parallelization should increase speed on MNIST-10k dataset"
+
+    def test_permutation():
+        X = np.loadtxt(os.path.join("data", "MNIST_10k.csv"))
+        kmed = KMedoids(n_medoids=5, algorithm="BanditPAM")
+        time_perm = time_measured_fit(kmed=kmed, X=X, loss="L2")
+        kmed.use_perm = False
+        time_no_perm = time_measured_fit(kmed=kmed, X=X, loss="L2")
+
+        print(time_perm, time_no_perm)
+        assert (
+            time_no_perm > time_perm
+        ), "Permutation should increase speed on MNIST-10k dataset"
+
+    def test_time_per_swap():
+        X = np.loadtxt(os.path.join("data", "MNIST_10k.csv"))
+        kmed = KMedoids(n_medoids=5, algorithm="BanditPAM")
+        time_ = time_measured_fit(kmed=kmed, X=X, loss="L2")
+        total_swap_time = kmed.total_swap_time
+        average_swap_time = kmed.time_per_swap
+
+        print(
+            time_,
+            total_swap_time / MILLISECONDS_IN_A_SECOND,
+            average_swap_time / MILLISECONDS_IN_A_SECOND,
+        )
+        assert (
+            time_ > total_swap_time / MILLISECONDS_IN_A_SECOND
+        ), "Total time should be greater than total swap time"
+        assert (
+            total_swap_time > average_swap_time / MILLISECONDS_IN_A_SECOND
+        ), "Total swap time should be greater than average swap time"
+        assert math.isclose(
+            total_swap_time / kmed.steps, average_swap_time, abs_tol=kmed.steps
+        ), "Average swap time inconsistent with total swap time and steps"
+
+    def test_old_bpam_vs_new_bpam():
+        X = np.loadtxt(os.path.join("data", "MNIST_1k.csv"))
+        kmed = KMedoids(
+            n_medoids=10,
+            algorithm="BanditPAM",
+            use_cache=False,
+            use_perm=False,
+            parallelize=False,
+        )
+        (
+            bpam_time,
+            bpam_width,
+            bpam_distance_computations,
+            bpam_steps,
+            bpam_writes,
+            bpam_hits,
+            bpam_misses,
+        ) = get_cache_statistics(kmed=kmed, X=X, loss="L2", cache_width=1000)
+
+        kmed_orig = KMedoids(
+            n_medoids=10,
+            algorithm="BanditPAM_orig",
+            use_cache=False,
+            use_perm=False,
+            parallelize=False,
+        )
+        (
+            bpam_orig_time,
+            bpam_orig_width,
+            bpam_orig_distance_computations,
+            bpam_orig_steps,
+            bpam_orig_writes,
+            bpam_orig_hits,
+            bpam_orig_misses,
+        ) = get_cache_statistics(
+            kmed=kmed_orig,
+            X=X,
+            loss="L2",
+            cache_width=1000,
+        )
+
+        print(bpam_time, bpam_orig_time)
+        print(kmed.time_per_swap, kmed_orig.time_per_swap)
+
+    # test_cache_stats()
+    # test_parallelization()
+    # test_permutation()
+    # test_time_per_swap()
+    test_old_bpam_vs_new_bpam()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/colab_files/colab_install_armadillo.sh b/scripts/colab_files/colab_install_armadillo.sh
index 6f0964eb..3d465681 100755
--- a/scripts/colab_files/colab_install_armadillo.sh
+++ b/scripts/colab_files/colab_install_armadillo.sh
@@ -2,7 +2,7 @@
 
 # Install dependencies on Google Colab machine
 
-apt install -y build-essential checkinstall libreadline-gplv2-dev libncursesw5-dev libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev
+apt install -y build-essential checkinstall libncursesw5-dev libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev
 
 # Copies and unzips the files to their appropriate places
 
diff --git a/scripts/colab_files/colab_setup.sh b/scripts/colab_files/colab_setup.sh
index b1cf2d96..2d4966b5 100755
--- a/scripts/colab_files/colab_setup.sh
+++ b/scripts/colab_files/colab_setup.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 
 apt-get update
-apt install build-essential checkinstall libreadline-gplv2-dev libncursesw5-dev libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev 
+apt install build-essential checkinstall libncursesw5-dev libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev
 
 # Should be in /content/
 cd /content
diff --git a/scripts/compare_banditpam_versions.py b/scripts/compare_banditpam_versions.py
new file mode 100644
index 00000000..252e5e94
--- /dev/null
+++ b/scripts/compare_banditpam_versions.py
@@ -0,0 +1,41 @@
+import numpy as np
+import os
+import banditpam
+import time
+
+import comparison_utils
+
+
+def main():
+    X1k = np.loadtxt(os.path.join("data", "MNIST_1k.csv"))
+    X10k = np.loadtxt(os.path.join("data", "MNIST_10k.csv"))
+    # X70k = np.loadtxt(os.path.join("data", "MNIST_70k.csv"))
+
+    for X in [X1k, X10k]:  # , X10k, X70k]:
+        new_kmed = banditpam.KMedoids(
+            n_medoids=10,
+            algorithm="BanditPAM",
+            use_cache=False,
+            use_perm=False,
+            parallelize=False,
+        )
+        start = time.time()
+        new_kmed.fit(X, "L2")
+        end = time.time()
+        comparison_utils.print_results(new_kmed, end - start)
+
+        old_kmed = banditpam.KMedoids(
+            n_medoids=10,
+            algorithm="BanditPAM_orig",
+            use_cache=False,
+            use_perm=False,
+            parallelize=False,
+        )
+        start = time.time()
+        old_kmed.fit(X, "L2")
+        end = time.time()
+        comparison_utils.print_results(old_kmed, end - start)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/comparison_utils.py b/scripts/comparison_utils.py
new file mode 100644
index 00000000..a9a9a9ed
--- /dev/null
+++ b/scripts/comparison_utils.py
@@ -0,0 +1,23 @@
+def print_results(kmed, runtime):
+    print("-----Results-----")
+    print("Algorithm:", kmed.algorithm)
+    print("Final Medoids:", kmed.medoids)
+    print("Loss:", kmed.average_loss)
+    print("Misc complexity:", f"{kmed.misc_distance_computations:,}")
+    print("Build complexity:", f"{kmed.build_distance_computations:,}")
+    print("Swap complexity:", f"{kmed.swap_distance_computations:,}")
+    print("Number of Swaps", kmed.steps)
+    print(
+        "Average Swap Sample Complexity:",
+        f"{kmed.swap_distance_computations / kmed.steps:,}",
+    )
+    print(
+        "Total complexity (without misc):",
+        f"{kmed.getDistanceComputations(False):,}"
+    )
+    print(
+        "Total complexity (with misc):",
+        f"{kmed.getDistanceComputations(True):,}",
+    )
+    print("Runtime per swap:", runtime / kmed.steps)
+    print("Total runtime:", runtime)
diff --git a/scripts/comparison_with_fasterpam.py b/scripts/comparison_with_fasterpam.py
new file mode 100644
index 00000000..b324e155
--- /dev/null
+++ b/scripts/comparison_with_fasterpam.py
@@ -0,0 +1,100 @@
+# Primarily authored by Erich Schubert (@kno10)
+
+from sklearn.datasets import fetch_openml
+from collections import defaultdict
+import time
+import numpy as np
+import banditpam
+
+from sklearn.metrics.pairwise import euclidean_distances
+from kmedoids import fasterpam
+
+# - Why is there a variable XX in the Google colab?
+
+
+def benchmark(data, f, n=1):
+    data_ = defaultdict(list)
+    for i in range(n):
+        print("Restart", i)
+        for k, v in f(data, i).items():
+            data_[k].append(v)
+
+    for k, v in data_.items():
+        v = np.array(v)
+        min, avg = v.min(), v.mean()
+        ste = v.std(ddof=1) / np.sqrt(len(v)) if len(v) > 1 else 0.0
+        print("{:16s} min={:-10.2f} mean={:-10.2f} ±{:-.2f}"
+              .format(k, min, avg, ste))
+
+
+def run_fasterpam(data, seed):
+    diss = euclidean_distances(data)
+    start = time.time()
+    r = fasterpam(diss, 5, random_state=seed)
+    end = time.time()
+    print("FasterPAM took ", r.n_iter, " iterations")
+    meds, lbl = data[r.medoids], r.labels
+    verified_loss = np.sqrt(((data - meds[lbl]) ** 2).sum(axis=1)).sum()
+    return {
+        "time (ms)": (end - start),
+        "verified loss": verified_loss,
+        "reported loss": r.loss,
+    }
+
+
+def run_bandit(data, seed):
+    diss = euclidean_distances(data)
+    km = banditpam.KMedoids(5, parallelize=True, dist_mat=diss)
+    print(km.algorithm)
+    km.seed = seed
+    start = time.time()
+    km.fit(data, "L2")
+    end = time.time()
+    meds, lbl = data[km.medoids], km.labels
+    verified_loss = np.sqrt(((data - meds[lbl]) ** 2).sum(axis=1)).sum()
+    return {
+        "time (ms)": (end - start),
+        "verified loss": verified_loss,
+        "reported loss": km.average_loss,
+    }
+
+
+def run_old_bandit(data, seed):
+    diss = euclidean_distances(data)
+    km = banditpam.KMedoids(
+        n_medoids=5,
+        parallelize=True,
+        algorithm="BanditPAM_orig",
+        dist_mat=diss)
+    print(km.algorithm)
+    km.seed = seed
+    start = time.time()
+    km.fit(data, "L2")
+    end = time.time()
+    meds, lbl = data[km.medoids], km.labels
+    verified_loss = np.sqrt(((data - meds[lbl]) ** 2).sum(axis=1)).sum()
+    return {
+        "time (ms)": (end - start),
+        "verified loss": verified_loss,
+        "reported loss": km.average_loss,
+    }
+
+
+if __name__ == "__main__":
+    X, _ = fetch_openml(
+        "mnist_784",
+        version=1,
+        return_X_y=True,
+        as_frame=False,
+    )
+    X = X[:20000]  # at 20k, colab will timeout for BanditPAM
+    print(X.shape, type(X))
+
+    # print("Benchmarking FasterPAM")
+    # benchmark(X, run_fasterpam)
+
+    print("\n\nBenchmarking BanditPAM")
+    benchmark(X, run_bandit)
+
+    print("\n\nBenchmarking Old BanditPAM")
+    benchmark(X, run_old_bandit)
diff --git a/scripts/kmeans_timing.py b/scripts/kmeans_timing.py
index 9965498f..d25803df 100644
--- a/scripts/kmeans_timing.py
+++ b/scripts/kmeans_timing.py
@@ -2,14 +2,14 @@
 import numpy as np
 from sklearn.cluster import KMeans
 
-X = np.loadtxt('data/MNIST_1k.csv')
+X = np.loadtxt("data/MNIST_1k.csv")
 kmeans = KMeans(
     n_clusters=5,
-    init='random',
+    init="random",
     n_init=1,
     copy_x=True,
-    algorithm='full',
-    )
+    algorithm="full",
+)
 start = time.time()
 kmeans.fit(X)
 print(time.time() - start, "seconds")
diff --git a/scripts/sample_complexity_with_k.py b/scripts/sample_complexity_with_k.py
new file mode 100644
index 00000000..d492d83e
--- /dev/null
+++ b/scripts/sample_complexity_with_k.py
@@ -0,0 +1,25 @@
+import numpy as np
+import os
+
+import banditpam
+import comparison_utils
+
+
+def main():
+    banditpam.set_num_threads(1)  # Turn off parallelization
+    X = np.loadtxt(os.path.join("data", "MNIST_10k.csv"))
+
+    for k in [5, 10, 20, 40, 80]:
+        kmed = banditpam.KMedoids(
+            n_medoids=k,
+            algorithm="BanditPAM",
+            use_cache=False,
+            use_perm=False,
+            parallelize=False,
+        )
+        kmed.fit(X, "L2")
+        comparison_utils.print_results(kmed)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/scaling_with_k.py b/scripts/scaling_with_k.py
new file mode 100755
index 00000000..07ef9ec6
--- /dev/null
+++ b/scripts/scaling_with_k.py
@@ -0,0 +1,31 @@
+# Use this script to verify BanditPAM is multithreaded
+# Should complete in <= 3 seconds
+import time
+import numpy as np
+import banditpam
+
+banditpam.set_num_threads(1)
+
+X = np.loadtxt("data/MNIST_1k.csv")
+
+for k in [5, 10, 30]:
+    kmed = banditpam.KMedoids(
+        n_medoids=k,
+        algorithm="BanditPAM",
+        use_cache=False,
+        use_perm=False,
+        parallelize=False,
+    )
+    start = time.time()
+    kmed.fit(X, "L2")
+    print(time.time() - start, "seconds")
+    print("Number of SWAP steps:", kmed.steps)
+    print(kmed.medoids)
+
+# for k in [5, 10, 30]:
+#     kmed = banditpam.KMedoids(n_medoids=k, algorithm="BanditPAM_orig")
+#     start = time.time()
+#     kmed.fit(X, "L2")
+#     print(time.time() - start, "seconds")
+#     print("Number of SWAP steps:", kmed.steps)
+#     print(kmed.medoids)
diff --git a/scripts/timing.py b/scripts/timing.py
index 2655f1d6..85ce4299 100755
--- a/scripts/timing.py
+++ b/scripts/timing.py
@@ -1,15 +1,15 @@
 # Use this script to verify BanditPAM is multithreaded
 # Should complete in <= 3 seconds
-
 import time
 import numpy as np
 import banditpam
 
-# banditpam.set_num_threads(1)
+banditpam.set_num_threads(1)
 
-X = np.loadtxt('data/MNIST_1k.csv')
+X = np.loadtxt("data/MNIST_1k.csv")
 kmed = banditpam.KMedoids(n_medoids=5, algorithm="BanditPAM")
 start = time.time()
 kmed.fit(X, "L2")
 print(time.time() - start, "seconds")
+print("Number of SWAP steps:", kmed.steps)
 print(kmed.medoids)
diff --git a/scripts/timing_dist_mat.py b/scripts/timing_dist_mat.py
new file mode 100755
index 00000000..a38e3a85
--- /dev/null
+++ b/scripts/timing_dist_mat.py
@@ -0,0 +1,20 @@
+# Use this script to verify BanditPAM is multithreaded
+# Should complete in <= 3 seconds
+
+import time
+import numpy as np
+import banditpam
+from sklearn.metrics.pairwise import euclidean_distances
+import os
+
+
+# banditpam.set_num_threads(1)
+
+X = np.loadtxt(os.path.join("data", "MNIST_10k.csv"))
+kmed = banditpam.KMedoids(n_medoids=5, algorithm="BanditPAM")
+start = time.time()
+diss = euclidean_distances(X)
+kmed.fit(X, "L2", dist_mat=diss)
+print(time.time() - start, "seconds")
+print("Number of SWAP steps:", kmed.steps)
+print(kmed.medoids)
diff --git a/setup.py b/setup.py
index 50a4449c..52911ba6 100644
--- a/setup.py
+++ b/setup.py
@@ -8,10 +8,10 @@
 import distutils.sysconfig
 import distutils.spawn
 
-__version__ = "3.0.4"
+__version__ = "4.0.0"
 
 # TODO(@motiwari): Move this to a separate file
-GHA = 'GITHUB_ACTIONS'
+GHA = "GITHUB_ACTIONS"
 
 
 class get_pybind_include(object):
@@ -21,8 +21,10 @@ class get_pybind_include(object):
     until it is actually installed via setup's setup_requires arg,
     so that the ``get_include()`` method can be invoked.
     """
+
     def __str__(self):
         import pybind11
+
         return pybind11.get_include()
 
 
@@ -33,8 +35,10 @@ class get_numpy_include(object):
     until it is actually installed via setup's setup_requires arg,
     so that the ``get_include()`` method can be invoked.
     """
+
     def __str__(self):
         import numpy
+
         return numpy.get_include()
 
 
@@ -45,17 +49,19 @@ def compiler_check():
     another one!
     """
     try:
-        return 'clang' \
-            if 'clang' in distutils.sysconfig.get_config_vars()["CC"] \
-            else 'gcc'
+        return (
+            "clang"
+            if "clang" in distutils.sysconfig.get_config_vars()["CC"]
+            else "gcc"
+        )
     except KeyError:
         # The 'CC' environment variable hasn't been set.
         # In this case, search for clang and gcc
         # Borrowed from https://github.com/clab/dynet/blob/master/setup.py
-        if distutils.spawn.find_executable('clang') is not None:
-            return 'clang'
-        elif distutils.spawn.find_executable('gcc') is not None:
-            return 'gcc'
+        if distutils.spawn.find_executable("clang") is not None:
+            return "clang"
+        elif distutils.spawn.find_executable("gcc") is not None:
+            return "gcc"
 
     raise Exception("No C++ compiler was found. Please install LLVM clang.")
 
@@ -86,18 +92,20 @@ def cpp_flag(compiler: str):
     The newer version is prefered over c++11 (when it is available).
     """
     compiler_name = compiler_check()
-    if compiler_name == 'clang':
+    if compiler_name == "clang":
         flags = ["-std=c++17", "-std=c++14", "-std=c++11"]
     else:
         # Assume gcc
-        flags = ["-std=c++1y"]
+        flags = ["-std=c++17"]
 
     for flag in flags:
         if has_flag(compiler, flag):
             return flag
 
-    raise RuntimeError("Unsupported compiler -- \
-        at least C++11 support is needed!")
+    raise RuntimeError(
+        "Unsupported compiler -- \
+        at least C++11 support is needed!"
+    )
 
 
 def check_brew_package(pkg_name: str):
@@ -118,8 +126,10 @@ def check_brew_installation():
     process = subprocess.Popen(brew_cmd, stdout=subprocess.PIPE)
     output, _error = process.communicate()
     if output.decode() == "":
-        raise Exception("Error: Need to install homebrew! \
-            Please see https://brew.sh")
+        raise Exception(
+            "Error: Need to install homebrew! \
+            Please see https://brew.sh"
+        )
 
 
 def check_numpy_installation():
@@ -157,8 +167,9 @@ def install_check_mac():
         )  # We need to use LLVM clang since Apple's doesn't support OpenMP
 
         # Set compiler to LLVM clang on Mac for OpenMP support
-        distutils.sysconfig.get_config_vars()["CC"] = \
-            os.path.join(llvm_loc, "bin", "clang")
+        distutils.sysconfig.get_config_vars()["CC"] = os.path.join(
+            llvm_loc, "bin", "clang"
+        )
 
 
 def check_omp_install_linux():
@@ -171,10 +182,10 @@ def check_armadillo_install_linux():
     # Since armadillo is a C++ extension, just check if it exists
     cmd = ["find", "/", "-iname", "armadillo"]
     process = subprocess.Popen(
-            cmd,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.DEVNULL,
-                            )
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.DEVNULL,
+    )
     output, _error = process.communicate()
     if output.decode() == "":
         print(
@@ -200,7 +211,7 @@ def check_linux_package_installation(pkg_name: str):
             "Error: Need to install %s! \
             Please ensure all dependencies are installed \
             via your package manager (apt, yum, etc.): \
-            build-essential checkinstall libreadline-gplv2-dev \
+            build-essential checkinstall \
             libncursesw5-dev libssl-dev libsqlite3-dev tk-dev \
             libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev"
             % (pkg_name)
@@ -211,11 +222,11 @@ def check_linux_package_installation(pkg_name: str):
 def install_ubuntu_pkgs():
     # TODO: Remove dangerous os.system() calls
     # See https://stackoverflow.com/a/51329156
-    os.system('apt update')
-    os.system('apt install -y \
+    os.system("apt update")
+    os.system(
+        "apt install -y \
         build-essential \
         checkinstall \
-        libreadline-gplv2-dev \
         libncursesw5-dev \
         libssl-dev \
         libsqlite3-dev \
@@ -223,7 +234,7 @@ def install_ubuntu_pkgs():
         libc6-dev \
         libbz2-dev \
         libffi-dev \
-        zlib1g-dev'
+        zlib1g-dev"
     )  # noqa: E124
 
 
@@ -233,6 +244,7 @@ def setup_colab(delete_source=False):
     # NOTE: This only works for Colab instances with Ubuntu 18.04 Runtimes!
     try:
         import google.colab  # noqa: F401
+
         in_colab = True
     except ModuleNotFoundError:
         in_colab = False
@@ -246,12 +258,20 @@ def setup_colab(delete_source=False):
         # and set delete_source=True always
         repo_location = os.path.join("/", "content", "BanditPAM")
         # Note the space after the git URL to separate the source and target
-        os.system('git clone https://github.com/ThrunGroup/BanditPAM.git ' +
-                  repo_location)
-        os.system(repo_location +
-                  '/scripts/colab_files/colab_install_armadillo.sh')
+        os.system(
+            "git clone https://github.com/ThrunGroup/BanditPAM.git "
+            + repo_location
+        )
+        os.system(
+            repo_location +
+            os.path.join(
+                "scripts",
+                "colab_files",
+                "colab_install_armadillo.sh",
+            )
+        )
         if delete_source:
-            os.system('rm -rf ' + repo_location)
+            os.system("rm -rf " + repo_location)
 
 
 def setup_paperspace_gradient():
@@ -263,15 +283,19 @@ def setup_paperspace_gradient():
     # TODO: Remove dangerous os.system() calls
     # See https://stackoverflow.com/a/51329156
 
-    in_paperspace_gradient = os.path.exists('/notebooks')
+    in_paperspace_gradient = os.path.exists("/notebooks")
 
     if in_paperspace_gradient:
         install_ubuntu_pkgs()
-        os.system('DEBIAN_FRONTEND=noninteractive TZ=America/NewYork \
-            apt install -y tk-dev')
-        os.system('git clone \
-            https://gitlab.com/conradsnicta/armadillo-code.git')
-        os.system('cd armadillo-code && cmake . && make install')
+        os.system(
+            "DEBIAN_FRONTEND=noninteractive TZ=America/NewYork \
+            apt install -y tk-dev"
+        )
+        os.system(
+            "git clone \
+            https://gitlab.com/conradsnicta/armadillo-code.git"
+        )
+        os.system("cd armadillo-code && cmake . && make install")
 
 
 def install_check_ubuntu():
@@ -279,7 +303,6 @@ def install_check_ubuntu():
     dependencies = [
         "build-essential",
         "checkinstall",
-        "libreadline-gplv2-dev",
         "libncursesw5-dev",
         "libssl-dev",
         "libsqlite3-dev",
@@ -327,8 +350,9 @@ class BuildExt(build_ext):
         install_check_mac()
         # Verify that we're either compiling with clang or
         # inside a Github Action
-        assert compiler_check() == 'clang' or os.environ.get(GHA, False), \
-            "Need to install LLVM clang!"
+        assert compiler_check() == "clang" or os.environ.get(
+            GHA, False
+        ), "Need to install LLVM clang!"
         darwin_opts = ["-stdlib=libc++", "-mmacosx-version-min=10.14", "-O3"]
         c_opts["unix"] += darwin_opts
         l_opts["unix"] += darwin_opts
@@ -349,16 +373,26 @@ def build_extensions(self):
         # TODO(@motiwari): on Windows, these flags are unrecognized
         opts.append(cpp_flag(self.compiler))
         opts.append("-O3")
-        if sys.platform == 'darwin' and os.environ.get(GHA, False):
-            opts.append('-Xpreprocessor -fopenmp')
+        if sys.platform == "darwin" and os.environ.get(GHA, False):
+            opts.append('-Xpreprocessor')  # NEEDS TO BE WITH NEXT LINE
+            opts.append('-fopenmp')  # NEEDS TO BE WITH PREVIOUS LINE
+
+            opts.append('-lomp')  # Potentially unused?
+            opts.append('-I/usr/local/opt/libomp/include')
+            opts.append('-L/usr/local/opt/libomp/lib')  # Potentially unused?
         else:
-            opts.append('-fopenmp')
+            opts.append("-fopenmp")
 
         compiler_name = compiler_check()
-        if sys.platform == "darwin" or compiler_name == 'clang':
-            link_opts.append('-lomp')
+        if (sys.platform == "darwin" and os.environ.get(GHA, False)):
+            link_opts.append('-lomp')  # Potentially unused?
+            link_opts.append('-I/usr/local/opt/libomp/include')
+            link_opts.append('-L/usr/local/opt/libomp/lib')  # Potentially unused?
         else:
-            link_opts.append("-lgomp")
+            if compiler_name == "clang":
+                link_opts.append("-lomp")
+            else:  # gcc
+                link_opts.append("-lgomp")
 
         if ct == "unix":
             if has_flag(self.compiler, "-fvisibility=hidden"):
@@ -366,10 +400,9 @@ def build_extensions(self):
 
         for ext in self.extensions:
             ext.define_macros = [
-                (
-                    "VERSION_INFO",
-                    '"{}"'.format(self.distribution.get_version())
-                )
+                ("VERSION_INFO", '"{}"'.format(
+                    self.distribution.get_version()
+                ))
             ]
             ext.extra_compile_args = opts
             ext.extra_link_args = link_opts
@@ -403,10 +436,10 @@ def main():
             # To include carma when the BanditPAM repo hasnt been initialized
             os.path.join("/", "usr", "local", "include"),
             os.path.join("/", "usr", "local", "include", "carma"),
-            os.path.join("/", "usr", "local", "include", "carma", \
+            os.path.join("/", "usr", "local", "include", "carma",
                          "carma_bits"),
             # When building from source on M1 Macs, may need these dirs
-            # Currently, we should never be building from soure on an M1 Mac,
+            # Currently, we should never be building from source on an M1 Mac,
             # Only cross-compiling from an Intel Mac
             # TODO(@motiwari): Remove extraneous directories
             os.path.join("/", "opt", "homebrew"),
@@ -415,18 +448,30 @@ def main():
             os.path.join("/", "opt", "homebrew", "lib"),
             os.path.join("/", "opt", "homebrew", "opt"),
             os.path.join("/", "opt", "homebrew", "opt", "armadillo"),
-            os.path.join("/", "opt", "homebrew", "opt", "armadillo", \
+            os.path.join("/", "opt", "homebrew", "opt", "armadillo",
                          "include"),
-            os.path.join("/", "opt", "homebrew", "opt", "armadillo", \
-                         "include", "armadillo_bits"),
+            os.path.join(
+                "/", "opt", "homebrew", "opt", "armadillo", "include",
+                "armadillo_bits"
+            ),
+            # Needed for Mac Github Runners
+            # for macos-10.15
+            os.path.join("/", "usr", "local", "Cellar", "libomp",
+                         "15.0.2", "include"),
+            # for macos-latest
+            os.path.join("/", "usr", "local", "Cellar", "libomp",
+                         "15.0.7", "include"),
         ]
 
     compiler_name = compiler_check()
-    if compiler_name == "clang" or \
-            (os.environ.get(GHA, False) and sys.platform == 'darwin'):
-        libraries = ["armadillo", "omp"]  # For M1 Mac runner build
-    else:  # gcc
-        libraries = ["armadillo", "gomp"]
+    if (sys.platform == "darwin" and os.environ.get(GHA, False)):
+        # On Mac Github Runners, we should NOT include gomp or omp here due to build errors.
+        libraries = ["armadillo", "omp"]
+    else:
+        if compiler_name == "clang":
+            libraries = ["armadillo", "omp"]
+        else:  # gcc
+            libraries = ["armadillo", "gomp"]
 
     ext_modules = [
         Extension(
@@ -435,21 +480,35 @@ def main():
                 os.path.join("src", "algorithms", "kmedoids_algorithm.cpp"),
                 os.path.join("src", "algorithms", "pam.cpp"),
                 os.path.join("src", "algorithms", "banditpam.cpp"),
+                os.path.join("src", "algorithms", "banditpam_orig.cpp"),
                 os.path.join("src", "algorithms", "fastpam1.cpp"),
                 os.path.join("src", "python_bindings",
                              "kmedoids_pywrapper.cpp"),
-                os.path.join("src", "python_bindings", "fit_python.cpp"),
                 os.path.join("src", "python_bindings", "medoids_python.cpp"),
                 os.path.join("src", "python_bindings",
                              "build_medoids_python.cpp"),
+                os.path.join("src", "python_bindings", "fit_python.cpp"),
                 os.path.join("src", "python_bindings", "labels_python.cpp"),
                 os.path.join("src", "python_bindings", "steps_python.cpp"),
                 os.path.join("src", "python_bindings", "loss_python.cpp"),
+                os.path.join("src", "python_bindings", "cache_python.cpp"),
+                os.path.join("src", "python_bindings",
+                             "swap_times_python.cpp"),
             ],
             include_dirs=include_dirs,
-            library_dirs=[os.path.join("/", "usr", "local", "lib")],
+            library_dirs=[
+                os.path.join("/", "usr", "local", "lib"),
+                # For Mac Github runners that install locally (not build wheels) - testing
+                # for macos-10.15
+                os.path.join("/", "usr", "local", "Cellar", "libomp",
+                             "15.0.2", "lib"),
+
+                # for macos-latest
+                os.path.join("/", "usr", "local", "Cellar", "libomp",
+                             "15.0.7", "lib"),
+            ],
             libraries=libraries,
-            language="c++1y",  # TODO: modify this based on cpp_flag(compiler)
+            language="c++1z",  # TODO: modify this based on cpp_flag(compiler)
             extra_compile_args=["-static-libstdc++"],
         )
     ]
@@ -460,7 +519,7 @@ def main():
     setup(
         name="banditpam",
         version=__version__,
-        author="Mo Tiwari and James Mayclin",
+        author="Mo Tiwari",
         maintainer="Mo Tiwari",
         author_email="motiwari@stanford.edu",
         url="https://github.com/ThrunGroup/BanditPAM",
@@ -479,8 +538,7 @@ def main():
             "Operating System :: OS Independent",
         ],
         headers=[
-            os.path.join("headers", "algorithms",
-                         "kmedoids_algorithm.hpp"),
+            os.path.join("headers", "algorithms", "kmedoids_algorithm.hpp"),
             os.path.join("headers", "algorithms", "banditpam.hpp"),
             os.path.join("headers", "algorithms", "fastpam1.hpp"),
             os.path.join("headers", "algorithms", "pam.hpp"),
@@ -490,5 +548,5 @@ def main():
     )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index a273009a..e4db8da2 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -4,9 +4,16 @@ include_directories(${PROJECT_SOURCE_DIR}/headers/python_bindings)
 include_directories(${PROJECT_SOURCE_DIR}/headers/carma/include)
 include_directories(${PROJECT_SOURCE_DIR}/headers/carma/include/carma_bits)
 
+# For Mac Github Runner to find OpenMP -- potentially unnecessary after fixing -Xpreprocessor -fopenmp issue
+include_directories(/usr/local/Cellar/libomp/15.0.2/include)
+include_directories(/usr/local/Cellar/libomp/15.0.7/include)
+include_directories(/usr/local/Cellar/libomp)
+include_directories(/usr/local/opt/libomp)
+include_directories(/usr/local/opt)
+
 add_executable(BanditPAM main.cpp)
 
-add_library(BanditPAM_LIB algorithms/kmedoids_algorithm.cpp algorithms/pam.cpp algorithms/banditpam.cpp algorithms/fastpam1.cpp)
+add_library(BanditPAM_LIB algorithms/kmedoids_algorithm.cpp algorithms/pam.cpp algorithms/banditpam.cpp algorithms/banditpam_orig.cpp algorithms/fastpam1.cpp)
 target_link_libraries(BanditPAM PUBLIC BanditPAM_LIB)
 
 find_package(OpenMP REQUIRED)
diff --git a/src/algorithms/banditpam.cpp b/src/algorithms/banditpam.cpp
index 5668dc62..5db29393 100644
--- a/src/algorithms/banditpam.cpp
+++ b/src/algorithms/banditpam.cpp
@@ -12,22 +12,28 @@
 #include <cmath>
 
 namespace km {
-void BanditPAM::fitBanditPAM(const arma::fmat& inputData) {
+void BanditPAM::fitBanditPAM(
+  const arma::fmat& inputData,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat) {
   data = arma::trans(inputData);
 
-  if (this->useCacheP) {
+  // Note: even if we are using a distance matrix, we compute the permutation
+  // in the block below because it is used elsewhere in the call stack
+  // TODO(@motiwari): Remove need for data or permutation through when using
+  //  a distance matrix
+  if (this->useCache) {
     size_t n = data.n_cols;
-    size_t m = fmin(n, ceil(log10(data.n_cols) * cacheMultiplier));
+    size_t m = fmin(n, cacheWidth);
     cache = new float[n * m];
 
-    #pragma omp parallel for
+    #pragma omp parallel for if (this->parallelize)
     for (size_t idx = 0; idx < m*n; idx++) {
       cache[idx] = -1;  // TODO(@motiwari): need better value here
     }
 
     permutation = arma::randperm(n);
     permutationIdx = 0;
-    reindex = {};
+    reindex = {};  // TODO(@motiwari): Can this intialization be removed?
     // TODO(@motiwari): Can we parallelize this?
     for (size_t counter = 0; counter < m; counter++) {
       reindex[permutation[counter]] = counter;
@@ -36,24 +42,27 @@ void BanditPAM::fitBanditPAM(const arma::fmat& inputData) {
 
   arma::fmat medoidMatrix(data.n_rows, nMedoids);
   arma::urowvec medoidIndices(nMedoids);
-  BanditPAM::build(data, &medoidIndices, &medoidMatrix);
   steps = 0;
+  BanditPAM::build(data, distMat, &medoidIndices, &medoidMatrix);
 
   medoidIndicesBuild = medoidIndices;
   arma::urowvec assignments(data.n_cols);
-  BanditPAM::swap(data, &medoidIndices, &medoidMatrix, &assignments);
+
+  BanditPAM::swap(data, distMat, &medoidIndices, &medoidMatrix, &assignments);
+
   medoidIndicesFinal = medoidIndices;
   labels = assignments;
 }
 
 arma::frowvec BanditPAM::buildSigma(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   const arma::frowvec& bestDistances,
   const bool useAbsolute) {
   size_t N = data.n_cols;
   arma::uvec referencePoints;
   // TODO(@motiwari): Make this wraparound properly as
-  // last batch_size elements are dropped
+  //  last batch_size elements are dropped
   if (usePerm) {
     if ((permutationIdx + batchSize - 1) >= N) {
       permutationIdx = 0;
@@ -69,16 +78,18 @@ arma::frowvec BanditPAM::buildSigma(
 
   arma::fvec sample(batchSize);
   arma::frowvec updated_sigma(N);
-  #pragma omp parallel for
+  #pragma omp parallel for if (this->parallelize)
   for (size_t i = 0; i < N; i++) {
     for (size_t j = 0; j < batchSize; j++) {
-      float cost = KMedoids::cachedLoss(data, i, referencePoints(j));
+      // 0 for MISC
+      float cost =
+          KMedoids::cachedLoss(data, distMat, i, referencePoints(j), 0);
       if (useAbsolute) {
         sample(j) = cost;
       } else {
         sample(j) = cost < bestDistances(referencePoints(j))
                           ? cost : bestDistances(referencePoints(j));
-              sample(j) -= bestDistances(referencePoints(j));
+        sample(j) -= bestDistances(referencePoints(j));
       }
     }
     updated_sigma(i) = arma::stddev(sample);
@@ -88,6 +99,7 @@ arma::frowvec BanditPAM::buildSigma(
 
 arma::frowvec BanditPAM::buildTarget(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   const arma::uvec* target,
   const arma::frowvec* bestDistances,
   const bool useAbsolute,
@@ -97,10 +109,10 @@ arma::frowvec BanditPAM::buildTarget(
   if (exact > 0) {
     tmpBatchSize = N;
   }
-  arma::frowvec estimates(target->n_rows, arma::fill::zeros);
+  arma::frowvec results(target->n_rows, arma::fill::zeros);
   arma::uvec referencePoints;
   // TODO(@motiwari): Make this wraparound properly
-  // as last batch_size elements are dropped
+  //  as last batch_size elements are dropped
   if (usePerm) {
     if ((permutationIdx + tmpBatchSize - 1) >= N) {
       permutationIdx = 0;
@@ -114,12 +126,17 @@ arma::frowvec BanditPAM::buildTarget(
     referencePoints = arma::randperm(N, tmpBatchSize);
   }
 
-  #pragma omp parallel for
+  #pragma omp parallel for if (this->parallelize)
   for (size_t i = 0; i < target->n_rows; i++) {
     float total = 0;
     for (size_t j = 0; j < referencePoints.n_rows; j++) {
       float cost =
-        KMedoids::cachedLoss(data, (*target)(i), referencePoints(j));
+        KMedoids::cachedLoss(
+            data,
+            distMat,
+            (*target)(i),
+            referencePoints(j),
+            1);  // 1 for BUILD
       if (useAbsolute) {
         total += cost;
       } else {
@@ -128,19 +145,20 @@ arma::frowvec BanditPAM::buildTarget(
         total -= (*bestDistances)(referencePoints(j));
       }
     }
-     estimates(i) = total / tmpBatchSize;
+    results(i) = total / tmpBatchSize;
   }
-  return estimates;
+  return results;
 }
 
 void BanditPAM::build(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   arma::urowvec* medoidIndices,
   arma::fmat* medoids) {
   size_t N = data.n_cols;
   arma::frowvec N_mat(N);
   N_mat.fill(N);
-  size_t p = (buildConfidence * N);
+  size_t p = N;
   bool useAbsolute = true;
   arma::frowvec estimates(N, arma::fill::zeros);
   arma::frowvec bestDistances(N);
@@ -152,27 +170,27 @@ void BanditPAM::build(
   arma::frowvec numSamples(N, arma::fill::zeros);
   arma::frowvec exactMask(N, arma::fill::zeros);
 
-  // TODO(@motiwari): #pragma omp parallel for?
+  // TODO(@motiwari): #pragma omp parallel for if (this->parallelize)?
   for (size_t k = 0; k < nMedoids; k++) {
     // instantiate medoids one-by-one
     permutationIdx = 0;
-    size_t step_count = 0;
     candidates.fill(1);
     numSamples.fill(0);
     exactMask.fill(0);
     estimates.fill(0);
     // compute std dev amongst batch of reference points
-    sigma = buildSigma(data, bestDistances, useAbsolute);
+    sigma = buildSigma(data, distMat, bestDistances, useAbsolute);
 
     while (arma::sum(candidates) > precision) {
       // TODO(@motiwari): Do not need a matrix for this comparison,
-      // use broadcasting
+      //  use broadcasting
       arma::umat compute_exactly =
         ((numSamples + batchSize) >= N_mat) != exactMask;
       if (arma::accu(compute_exactly) > 0) {
         arma::uvec targets = find(compute_exactly);
         arma::frowvec result = buildTarget(
           data,
+          distMat,
           &targets,
           &bestDistances,
           useAbsolute,
@@ -190,6 +208,7 @@ void BanditPAM::build(
       arma::uvec targets = arma::find(candidates);
       arma::frowvec result = buildTarget(
         data,
+        distMat,
         &targets,
         &bestDistances,
         useAbsolute,
@@ -202,23 +221,28 @@ void BanditPAM::build(
       numSamples.cols(targets) += batchSize;
       arma::frowvec adjust(targets.n_rows);
       adjust.fill(p);
-      adjust = arma::log(adjust);
+      // Assume buildConfidence is given in logspace
+      adjust = buildConfidence + arma::log(adjust);
       arma::frowvec confBoundDelta =
         sigma.cols(targets) %
         arma::sqrt(adjust / numSamples.cols(targets));
       ucbs.cols(targets) = estimates.cols(targets) + confBoundDelta;
       lcbs.cols(targets) = estimates.cols(targets) - confBoundDelta;
       candidates = (lcbs < ucbs.min()) && (exactMask == 0);
-      step_count++;
     }
 
     medoidIndices->at(k) = lcbs.index_min();
     medoids->unsafe_col(k) = data.unsafe_col((*medoidIndices)(k));
 
     // don't need to do this on final iteration
-    #pragma omp parallel for
+    #pragma omp parallel for if (this->parallelize)
     for (size_t i = 0; i < N; i++) {
-        float cost = KMedoids::cachedLoss(data, i, (*medoidIndices)(k));
+        float cost = KMedoids::cachedLoss(
+          data,
+          distMat,
+          i,
+          (*medoidIndices)(k),
+          0);  // 0 for MISC
         if (cost < bestDistances(i)) {
             bestDistances(i) = cost;
         }
@@ -230,6 +254,7 @@ void BanditPAM::build(
 
 arma::fmat BanditPAM::swapSigma(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   const arma::frowvec* bestDistances,
   const arma::frowvec* secondBestDistances,
   const arma::urowvec* assignments) {
@@ -238,7 +263,7 @@ arma::fmat BanditPAM::swapSigma(
   arma::fmat updated_sigma(K, N, arma::fill::zeros);
   arma::uvec referencePoints;
   // TODO(@motiwari): Make this wraparound properly
-  // as last batch_size elements are dropped
+  //  as last batch_size elements are dropped
   if (usePerm) {
     if ((permutationIdx + batchSize - 1) >= N) {
       permutationIdx = 0;
@@ -254,7 +279,7 @@ arma::fmat BanditPAM::swapSigma(
 
   arma::fvec sample(batchSize);
   // for each considered swap
-  #pragma omp parallel for
+  #pragma omp parallel for if (this->parallelize)
   for (size_t i = 0; i < K * N; i++) {
     // extract data point of swap
     size_t n = i / K;
@@ -262,7 +287,9 @@ arma::fmat BanditPAM::swapSigma(
 
     // calculate change in loss for some subset of the data
     for (size_t j = 0; j < batchSize; j++) {
-      float cost = KMedoids::cachedLoss(data, n, referencePoints(j));
+      // 0 for MISC when estimating sigma
+      float cost =
+          KMedoids::cachedLoss(data, distMat, n, referencePoints(j), 0);
 
       if (k == (*assignments)(referencePoints(j))) {
         if (cost < (*secondBestDistances)(referencePoints(j))) {
@@ -284,16 +311,36 @@ arma::fmat BanditPAM::swapSigma(
   return updated_sigma;
 }
 
-arma::fvec BanditPAM::swapTarget(
+arma::fmat BanditPAM::swapTarget(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   const arma::urowvec* medoidIndices,
   const arma::uvec* targets,
   const arma::frowvec* bestDistances,
   const arma::frowvec* secondBestDistances,
   const arma::urowvec* assignments,
   const size_t exact = 0) {
-  size_t N = data.n_cols;
-  arma::fvec estimates(targets->n_rows, arma::fill::zeros);
+  const size_t N = data.n_cols;
+  const size_t T = targets->n_rows;
+  arma::fmat results(nMedoids, T, arma::fill::zeros);
+
+  // Targets should be a list of indices for target CANDIDATE points
+  // Then update all corresponding EXISTING MEDOID indices targets.
+  // If targets is a T-length vector, then the return value should be
+  // a matrix of size K x T. We should perform the appropriate update then
+  // in the swap() function.
+  //
+  // An alternate method to do this would be to pass only the (m, c)
+  // Points under consideration. Then we wouldn't need to update all
+  // k virtual arms for each candidate, just the ones that are passed
+  // However, this would incur a .find() call to find all pairs
+  // (m, c) where c == c', the arm under consideration. I believe this
+  // would be an O(kn) cost. Instead, may need to use another data
+  // structure to avoid this .find() call, like a tree where the top-level
+  // nodes are the candidates and the bottom-level nodes are the corresponding
+  // virtual arms.
+  // A jagged array might also do the trick.
+
 
   size_t tmpBatchSize = batchSize;
   if (exact > 0) {
@@ -302,7 +349,8 @@ arma::fvec BanditPAM::swapTarget(
 
   arma::uvec referencePoints;
   // TODO(@motiwari): Make this wraparound properly
-  // as last batch_size elements are dropped
+  //  as last batch_size elements are dropped
+  // TODO(@motiwari): Break this duplicated code into a function
   if (usePerm) {
     if ((permutationIdx + tmpBatchSize - 1) >= N) {
       permutationIdx = 0;
@@ -316,43 +364,47 @@ arma::fvec BanditPAM::swapTarget(
     referencePoints = arma::randperm(N, tmpBatchSize);
   }
 
-  // for each considered swap
-  #pragma omp parallel for
-  for (size_t i = 0; i < targets->n_rows; i++) {
-    float total = 0;
-    // extract data point of swap
-    size_t n = (*targets)(i) / medoidIndices->n_cols;
-    size_t k = (*targets)(i) % medoidIndices->n_cols;
-    // calculate total loss for some subset of the data
+  // TODO(@motiwari): Declare variables outside of loops
+  #pragma omp parallel for if (this->parallelize)
+  for (size_t i = 0; i < T; i++) {
+    // TODO(@motiwari): pragma omp parallel for?
     for (size_t j = 0; j < tmpBatchSize; j++) {
-      float cost = KMedoids::cachedLoss(data, n, referencePoints(j));
-      if (k == (*assignments)(referencePoints(j))) {
-        if (cost < (*secondBestDistances)(referencePoints(j))) {
-          total += cost;
-        } else {
-          total += (*secondBestDistances)(referencePoints(j));
-        }
-      } else {
-        if (cost < (*bestDistances)(referencePoints(j))) {
-          total += cost;
-        } else {
-          total += (*bestDistances)(referencePoints(j));
-        }
+      float cost =
+          KMedoids::cachedLoss(
+              data,
+              distMat,
+              (*targets)(i),
+              referencePoints(j),
+              2);  // 2 for SWAP
+      size_t k = (*assignments)(referencePoints(j));
+      if (cost < (*bestDistances)(referencePoints(j))) {
+          // We might be able to change this to .eachrow(every column but k)
+          // since arma does this in-place and it should not introduce
+          // complexity
+          results.col(i) += cost - (*bestDistances)(referencePoints(j));
       }
-      total -= (*bestDistances)(referencePoints(j));
+
+      // If cost < bd, this second term will subtract off the "new cost"
+      // added by the all-column call above inside the if
+      results(k, i) +=
+        std::fmin(cost, (*secondBestDistances)(referencePoints(j))) -
+          std::fmin(cost, (*bestDistances)(referencePoints(j)));
     }
-    estimates(i) = total / referencePoints.n_rows;
   }
-  return estimates;
+  // TODO(@motiwari): we can probably avoid this division
+  //  if we look at total loss, not average loss
+  results /= tmpBatchSize;
+  return results;
 }
 
 void BanditPAM::swap(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   arma::urowvec* medoidIndices,
   arma::fmat* medoids,
   arma::urowvec* assignments) {
   size_t N = data.n_cols;
-  size_t p = (N * nMedoids * swapConfidence);
+  size_t p = N;
 
   arma::fmat sigma(nMedoids, N, arma::fill::zeros);
 
@@ -369,6 +421,7 @@ void BanditPAM::swap(
   // calculate quantities needed for swap, bestDistances and sigma
   calcBestDistancesSwap(
     data,
+    distMat,
     medoidIndices,
     &bestDistances,
     &secondBestDistances,
@@ -382,87 +435,120 @@ void BanditPAM::swap(
 
     sigma = swapSigma(
       data,
+      distMat,
       &bestDistances,
       &secondBestDistances,
       assignments);
 
+    // Reset variables when starting a new swap
     candidates.fill(1);
     exactMask.fill(0);
     estimates.fill(0);
     numSamples.fill(0);
 
     // while there is at least one candidate (float comparison issues)
-    while (arma::accu(candidates) > 0.5) {
+    while (arma::accu(candidates) > 1.5) {
       // compute exactly if it's been samples more than N times and
       // hasn't been computed exactly already
       arma::umat compute_exactly =
         ((numSamples + batchSize) >= N) != (exactMask);
-      arma::uvec targets = arma::find(compute_exactly);
 
-      if (targets.size() > 0) {
-          arma::fvec result = swapTarget(
-            data,
-            medoidIndices,
-            &targets,
-            &bestDistances,
-            &secondBestDistances,
-            assignments,
-            N);
-          estimates.elem(targets) = result;
-          ucbs.elem(targets) = result;
-          lcbs.elem(targets) = result;
-          exactMask.elem(targets).fill(1);
-          numSamples.elem(targets) += N;
-          candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+      // Get unique candidate medoids from the candidates (second index)
+      // Store all k x T in estimates
+      // TODO(@motiwari): Move this declaration outside loop
+      // Need unique values over second index
+      // Sum the different columns
+      // if any index appears in at least one, compute it exactly
+      // TODO(@motiwari): make sure we're only computing exactly
+      // for the relevant candidates
+      arma::uvec compute_exactly_targets =
+        arma::find(arma::sum(compute_exactly, 0) >= 1);
+
+      if (compute_exactly_targets.size() > 0) {
+        arma::fmat result = swapTarget(
+          data,
+          distMat,
+          medoidIndices,
+          &compute_exactly_targets,
+          &bestDistances,
+          &secondBestDistances,
+          assignments,
+          N);
+
+        // result will be k x T
+        // Now update the correct indices
+        estimates.cols(compute_exactly_targets) = result;
+        ucbs.cols(compute_exactly_targets) = result;
+        lcbs.cols(compute_exactly_targets) = result;
+        exactMask.cols(compute_exactly_targets).fill(1);
+        numSamples.cols(compute_exactly_targets) += N;
+        candidates = (lcbs < ucbs.min()) && (exactMask == 0);
       }
       if (arma::accu(candidates) < precision) {
         break;
       }
-      targets = arma::find(candidates);
-      arma::fvec result = swapTarget(
+
+      // candidate_targets should be of size T
+      // Sum the different columns
+      // if any index appears in at least one column, sample it
+      arma::uvec candidate_targets = arma::find(arma::sum(candidates, 0) >= 1);
+
+      // result will be k x T
+      arma::fmat result = swapTarget(
         data,
+        distMat,
         medoidIndices,
-        &targets,
+        &candidate_targets,
         &bestDistances,
         &secondBestDistances,
         assignments,
         0);
-      estimates.elem(targets) =
-        ((numSamples.elem(targets) % estimates.elem(targets)) +
-        (result * batchSize)) /
-        (batchSize + numSamples.elem(targets));
-      numSamples.elem(targets) += batchSize;
-      arma::fvec adjust(targets.n_rows);
+
+      // candidate_targets should be of size T, 1
+      estimates.cols(candidate_targets) =
+        ((numSamples.cols(candidate_targets)
+        % estimates.cols(candidate_targets))
+        + (result * batchSize)) / (batchSize +
+          numSamples.cols(candidate_targets));
+
+      // numSamples should be k x N
+      // select the T of N columns that are candidates
+      numSamples.cols(candidate_targets) += batchSize;
+
+      arma::fmat adjust(nMedoids, candidate_targets.size());
+      // TODO(@motiwari): Move this ::fill to the previous line
       adjust.fill(p);
-      adjust = arma::log(adjust);
-      arma::fvec confBoundDelta = sigma.elem(targets) %
-                          arma::sqrt(adjust / numSamples.elem(targets));
+      // Assume swapConfidence is given in logspace
+      adjust = swapConfidence + arma::log(adjust);
+      arma::fmat confBoundDelta = sigma.cols(candidate_targets) %
+        arma::sqrt(adjust / numSamples.cols(candidate_targets));
+      ucbs.cols(candidate_targets) = estimates.cols(candidate_targets)
+        + confBoundDelta;
+      lcbs.cols(candidate_targets) = estimates.cols(candidate_targets)
+        - confBoundDelta;
 
-      ucbs.elem(targets) = estimates.elem(targets) + confBoundDelta;
-      lcbs.elem(targets) = estimates.elem(targets) - confBoundDelta;
       candidates = (lcbs < ucbs.min()) && (exactMask == 0);
-      targets = arma::find(candidates);
     }
 
     // Perform the medoid switch
     arma::uword newMedoid = lcbs.index_min();
-    // extract old and new medoids of swap
-    size_t k = newMedoid % medoids->n_cols;
-    size_t n = newMedoid / medoids->n_cols;
+    size_t k = newMedoid % nMedoids;
+    size_t n = newMedoid / nMedoids;
     swapPerformed = (*medoidIndices)(k) != n;
-    steps++;
 
     if (swapPerformed) {
       (*medoidIndices)(k) = n;
       medoids->col(k) = data.col((*medoidIndices)(k));
     }
+
     calcBestDistancesSwap(
-        data,
-        medoidIndices,
-        &bestDistances,
-        &secondBestDistances,
-        assignments,
-        swapPerformed);
+      data,
+      distMat,
+      medoidIndices,
+      &bestDistances,
+      &secondBestDistances,
+      assignments,
+      swapPerformed);
   }
 }
 }  // namespace km
diff --git a/src/algorithms/banditpam_orig.cpp b/src/algorithms/banditpam_orig.cpp
new file mode 100644
index 00000000..7caef1d2
--- /dev/null
+++ b/src/algorithms/banditpam_orig.cpp
@@ -0,0 +1,512 @@
+/**
+ * @file banditpam_orig.cpp
+ * @date 2021-07-25
+ *
+ * Contains the primary C++ implementation of the BanditPAM_orig code.
+ */
+
+#include "banditpam_orig.hpp"
+
+#include <armadillo>
+#include <unordered_map>
+#include <cmath>
+
+namespace km {
+void BanditPAM_orig::fitBanditPAM_orig(
+  const arma::fmat& inputData,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat) {
+  data = arma::trans(inputData);
+
+  // Note: even if we are using a distance matrix, we compute the permutation
+  // in the block below because it is used elsewhere in the call stack
+  // TODO(@motiwari): Remove need for data or permutation through when using
+  //  a distance matrix
+  if (this->useCache) {
+    size_t n = data.n_cols;
+    size_t m = fmin(n, cacheWidth);
+    cache = new float[n * m];
+
+    #pragma omp parallel for if (this->parallelize)
+    for (size_t idx = 0; idx < m*n; idx++) {
+      cache[idx] = -1;  // TODO(@motiwari): need better value here
+    }
+
+    permutation = arma::randperm(n);
+    permutationIdx = 0;
+    reindex = {};  // TODO(@motiwari): Can this intialization be removed?
+    // TODO(@motiwari): Can we parallelize this?
+    for (size_t counter = 0; counter < m; counter++) {
+      reindex[permutation[counter]] = counter;
+    }
+}
+
+  arma::fmat medoidMatrix(data.n_rows, nMedoids);
+  arma::urowvec medoidIndices(nMedoids);
+  steps = 0;
+  BanditPAM_orig::build(data, distMat, &medoidIndices, &medoidMatrix);
+
+  medoidIndicesBuild = medoidIndices;
+  arma::urowvec assignments(data.n_cols);
+
+  BanditPAM_orig::swap(
+      data,
+      distMat,
+      &medoidIndices,
+      &medoidMatrix,
+      &assignments);
+
+  medoidIndicesFinal = medoidIndices;
+  labels = assignments;
+}
+
+arma::frowvec BanditPAM_orig::buildSigma(
+  const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+  const arma::frowvec& bestDistances,
+  const bool useAbsolute) {
+  size_t N = data.n_cols;
+  arma::uvec referencePoints;
+  // TODO(@motiwari): Make this wraparound properly as
+  //  last batch_size elements are dropped
+  if (usePerm) {
+    if ((permutationIdx + batchSize - 1) >= N) {
+      permutationIdx = 0;
+    }
+    // inclusive of both indices
+    referencePoints = permutation.subvec(
+      permutationIdx,
+      permutationIdx + batchSize - 1);
+    permutationIdx += batchSize;
+  } else {
+    referencePoints = arma::randperm(N, batchSize);
+  }
+
+  arma::fvec sample(batchSize);
+  arma::frowvec updated_sigma(N);
+  #pragma omp parallel for if (this->parallelize)
+  for (size_t i = 0; i < N; i++) {
+    for (size_t j = 0; j < batchSize; j++) {
+      // 0 for MISC
+      float cost =
+          KMedoids::cachedLoss(data, distMat, i, referencePoints(j), 0);
+      if (useAbsolute) {
+        sample(j) = cost;
+      } else {
+        sample(j) = cost < bestDistances(referencePoints(j))
+                          ? cost : bestDistances(referencePoints(j));
+        sample(j) -= bestDistances(referencePoints(j));
+      }
+    }
+    updated_sigma(i) = arma::stddev(sample);
+  }
+  return updated_sigma;
+}
+
+arma::frowvec BanditPAM_orig::buildTarget(
+  const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+  const arma::uvec* target,
+  const arma::frowvec* bestDistances,
+  const bool useAbsolute,
+  const size_t exact = 0) {
+  size_t N = data.n_cols;
+  size_t tmpBatchSize = batchSize;
+  if (exact > 0) {
+    tmpBatchSize = N;
+  }
+  arma::frowvec estimates(target->n_rows, arma::fill::zeros);
+  arma::uvec referencePoints;
+  // TODO(@motiwari): Make this wraparound properly
+  //  as last batch_size elements are dropped
+  if (usePerm) {
+    if ((permutationIdx + tmpBatchSize - 1) >= N) {
+      permutationIdx = 0;
+    }
+    // inclusive of both indices
+    referencePoints = permutation.subvec(
+      permutationIdx,
+      permutationIdx + tmpBatchSize - 1);
+    permutationIdx += tmpBatchSize;
+  } else {
+    referencePoints = arma::randperm(N, tmpBatchSize);
+  }
+
+  #pragma omp parallel for if (this->parallelize)
+  for (size_t i = 0; i < target->n_rows; i++) {
+    float total = 0;
+    for (size_t j = 0; j < referencePoints.n_rows; j++) {
+      float cost =
+          KMedoids::cachedLoss(
+              data,
+              distMat,
+              (*target)(i),
+              referencePoints(j),
+              1);  // 1 for BUILD
+      if (useAbsolute) {
+        total += cost;
+      } else {
+        total += cost < (*bestDistances)(referencePoints(j))
+                      ? cost : (*bestDistances)(referencePoints(j));
+        total -= (*bestDistances)(referencePoints(j));
+      }
+    }
+    estimates(i) = total / tmpBatchSize;
+  }
+  return estimates;
+}
+
+void BanditPAM_orig::build(
+  const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+  arma::urowvec* medoidIndices,
+  arma::fmat* medoids) {
+  size_t N = data.n_cols;
+  arma::frowvec N_mat(N);
+  N_mat.fill(N);
+  size_t p = N;
+  bool useAbsolute = true;
+  arma::frowvec estimates(N, arma::fill::zeros);
+  arma::frowvec bestDistances(N);
+  bestDistances.fill(std::numeric_limits<float>::infinity());
+  arma::frowvec sigma(N);
+  arma::urowvec candidates(N, arma::fill::ones);
+  arma::frowvec lcbs(N);
+  arma::frowvec ucbs(N);
+  arma::frowvec numSamples(N, arma::fill::zeros);
+  arma::frowvec exactMask(N, arma::fill::zeros);
+
+  // TODO(@motiwari): #pragma omp parallel for if (this->parallelize)?
+  for (size_t k = 0; k < nMedoids; k++) {
+    // instantiate medoids one-by-one
+    permutationIdx = 0;
+    candidates.fill(1);
+    numSamples.fill(0);
+    exactMask.fill(0);
+    estimates.fill(0);
+    // compute std dev amongst batch of reference points
+    sigma = buildSigma(data, distMat, bestDistances, useAbsolute);
+
+    while (arma::sum(candidates) > precision) {
+      // TODO(@motiwari): Do not need a matrix for this comparison,
+      //  use broadcasting
+      arma::umat compute_exactly =
+          ((numSamples + batchSize) >= N_mat) != exactMask;
+      if (arma::accu(compute_exactly) > 0) {
+        arma::uvec targets = find(compute_exactly);
+        arma::frowvec result = buildTarget(
+          data,
+          distMat,
+          &targets,
+          &bestDistances,
+          useAbsolute,
+          N);
+        estimates.cols(targets) = result;
+        ucbs.cols(targets) = result;
+        lcbs.cols(targets) = result;
+        exactMask.cols(targets).fill(1);
+        numSamples.cols(targets) += N;
+        candidates.cols(targets).fill(0);
+      }
+      if (arma::sum(candidates) < precision) {
+        break;
+      }
+      arma::uvec targets = arma::find(candidates);
+      arma::frowvec result = buildTarget(
+        data,
+        distMat,
+        &targets,
+        &bestDistances,
+        useAbsolute,
+        0);
+      // update the running average
+      estimates.cols(targets) =
+        ((numSamples.cols(targets) % estimates.cols(targets)) +
+         (result * batchSize)) /
+        (batchSize + numSamples.cols(targets));
+      numSamples.cols(targets) += batchSize;
+      arma::frowvec adjust(targets.n_rows);
+      adjust.fill(p);
+      // Assume buildConfidence is given in logspace
+      adjust = buildConfidence + arma::log(adjust);
+      arma::frowvec confBoundDelta =
+        sigma.cols(targets) %
+        arma::sqrt(adjust / numSamples.cols(targets));
+      ucbs.cols(targets) = estimates.cols(targets) + confBoundDelta;
+      lcbs.cols(targets) = estimates.cols(targets) - confBoundDelta;
+      candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+    }
+
+    medoidIndices->at(k) = lcbs.index_min();
+    medoids->unsafe_col(k) = data.unsafe_col((*medoidIndices)(k));
+
+    // don't need to do this on final iteration
+    #pragma omp parallel for if (this->parallelize)
+    for (size_t i = 0; i < N; i++) {
+      float cost = KMedoids::cachedLoss(
+        data,
+        distMat,
+        i,
+        (*medoidIndices)(k),
+        0);  // 0 for MISC
+      if (cost < bestDistances(i)) {
+        bestDistances(i) = cost;
+      }
+    }
+    // use difference of loss for sigma and sampling, not absolute
+    useAbsolute = false;
+  }
+}
+
+arma::fmat BanditPAM_orig::swapSigma(
+  const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+  const arma::frowvec* bestDistances,
+  const arma::frowvec* secondBestDistances,
+  const arma::urowvec* assignments) {
+  size_t N = data.n_cols;
+  size_t K = nMedoids;
+  arma::fmat updated_sigma(K, N, arma::fill::zeros);
+  arma::uvec referencePoints;
+  // TODO(@motiwari): Make this wraparound properly
+  //  as last batch_size elements are dropped
+  if (usePerm) {
+    if ((permutationIdx + batchSize - 1) >= N) {
+      permutationIdx = 0;
+    }
+    // inclusive of both indices
+    referencePoints = permutation.subvec(
+      permutationIdx,
+      permutationIdx + batchSize - 1);
+    permutationIdx += batchSize;
+  } else {
+    referencePoints = arma::randperm(N, batchSize);
+  }
+
+  arma::fvec sample(batchSize);
+  // for each considered swap
+  #pragma omp parallel for if (this->parallelize)
+  for (size_t i = 0; i < K * N; i++) {
+    // extract data point of swap
+    size_t n = i / K;
+    size_t k = i % K;
+
+    // calculate change in loss for some subset of the data
+    for (size_t j = 0; j < batchSize; j++) {
+      // 0 for MISC when estimating sigma
+      float cost =
+          KMedoids::cachedLoss(data, distMat, n, referencePoints(j), 0);
+
+      if (k == (*assignments)(referencePoints(j))) {
+        if (cost < (*secondBestDistances)(referencePoints(j))) {
+          sample(j) = cost;
+        } else {
+          sample(j) = (*secondBestDistances)(referencePoints(j));
+        }
+      } else {
+        if (cost < (*bestDistances)(referencePoints(j))) {
+          sample(j) = cost;
+        } else {
+          sample(j) = (*bestDistances)(referencePoints(j));
+        }
+      }
+      sample(j) -= (*bestDistances)(referencePoints(j));
+    }
+    updated_sigma(k, n) = arma::stddev(sample);
+  }
+  return updated_sigma;
+}
+
+arma::fvec BanditPAM_orig::swapTarget(
+  const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+  const arma::urowvec* medoidIndices,
+  const arma::uvec* targets,
+  const arma::frowvec* bestDistances,
+  const arma::frowvec* secondBestDistances,
+  const arma::urowvec* assignments,
+  const size_t exact = 0) {
+  size_t N = data.n_cols;
+  arma::fvec estimates(targets->n_rows, arma::fill::zeros);
+
+  size_t tmpBatchSize = batchSize;
+  if (exact > 0) {
+    tmpBatchSize = N;
+  }
+
+  arma::uvec referencePoints;
+  // TODO(@motiwari): Make this wraparound properly
+  //  as last batch_size elements are dropped
+  // TODO(@motiwari): Break this duplicated code into a function
+  if (usePerm) {
+    if ((permutationIdx + tmpBatchSize - 1) >= N) {
+      permutationIdx = 0;
+    }
+    // inclusive of both indices
+    referencePoints = permutation.subvec(
+      permutationIdx,
+      permutationIdx + tmpBatchSize - 1);
+    permutationIdx += tmpBatchSize;
+  } else {
+    referencePoints = arma::randperm(N, tmpBatchSize);
+  }
+
+  // TODO(@motiwari): Declare variables outside of loops
+  #pragma omp parallel for if (this->parallelize)
+  for (size_t i = 0; i < targets->n_rows; i++) {
+    float total = 0;
+    // extract data point of swap
+    size_t n = (*targets)(i) / medoidIndices->n_cols;
+    size_t k = (*targets)(i) % medoidIndices->n_cols;
+    // calculate total loss for some subset of the data
+    for (size_t j = 0; j < tmpBatchSize; j++) {
+      // 2 for SWAP
+      float cost =
+          KMedoids::cachedLoss(data, distMat, n, referencePoints(j), 2);
+      if (k == (*assignments)(referencePoints(j))) {
+        if (cost < (*secondBestDistances)(referencePoints(j))) {
+          total += cost;
+        } else {
+          total += (*secondBestDistances)(referencePoints(j));
+        }
+      } else {
+        if (cost < (*bestDistances)(referencePoints(j))) {
+          total += cost;
+        } else {
+          total += (*bestDistances)(referencePoints(j));
+        }
+      }
+      total -= (*bestDistances)(referencePoints(j));
+    }
+    // TODO(@motiwari): we can probably avoid this division
+    //  if we look at total loss, not average loss
+    estimates(i) = total / referencePoints.n_rows;
+  }
+  return estimates;
+}
+
+void BanditPAM_orig::swap(
+  const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
+  arma::urowvec* medoidIndices,
+  arma::fmat* medoids,
+  arma::urowvec* assignments) {
+  size_t N = data.n_cols;
+  size_t p = N;
+
+  arma::fmat sigma(nMedoids, N, arma::fill::zeros);
+
+  arma::frowvec bestDistances(N);
+  arma::frowvec secondBestDistances(N);
+  bool swapPerformed = true;
+  arma::umat candidates(nMedoids, N, arma::fill::ones);
+  arma::umat exactMask(nMedoids, N, arma::fill::zeros);
+  arma::fmat estimates(nMedoids, N, arma::fill::zeros);
+  arma::fmat lcbs(nMedoids, N);
+  arma::fmat ucbs(nMedoids, N);
+  arma::umat numSamples(nMedoids, N, arma::fill::zeros);
+
+  // calculate quantities needed for swap, bestDistances and sigma
+  calcBestDistancesSwap(
+    data,
+    distMat,
+    medoidIndices,
+    &bestDistances,
+    &secondBestDistances,
+    assignments,
+    swapPerformed);
+
+  // continue making swaps while loss is decreasing
+  while (swapPerformed && steps < maxIter) {
+    steps++;
+    permutationIdx = 0;
+
+    sigma = swapSigma(
+      data,
+      distMat,
+      &bestDistances,
+      &secondBestDistances,
+      assignments);
+
+    // Reset variables when starting a new swap
+    candidates.fill(1);
+    exactMask.fill(0);
+    estimates.fill(0);
+    numSamples.fill(0);
+
+    // while there is at least one candidate (float comparison issues)
+    while (arma::accu(candidates) > 1.5) {
+      // compute exactly if it's been samples more than N times and
+      // hasn't been computed exactly already
+      arma::umat compute_exactly =
+          ((numSamples + batchSize) >= N) != (exactMask);
+      arma::uvec targets = arma::find(compute_exactly);
+
+      if (targets.size() > 0) {
+        arma::fvec result = swapTarget(
+          data,
+          distMat,
+          medoidIndices,
+          &targets,
+          &bestDistances,
+          &secondBestDistances,
+          assignments,
+          N);
+        estimates.elem(targets) = result;
+        ucbs.elem(targets) = result;
+        lcbs.elem(targets) = result;
+        exactMask.elem(targets).fill(1);
+        numSamples.elem(targets) += N;
+        candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+      }
+      if (arma::accu(candidates) < precision) {
+        break;
+      }
+      targets = arma::find(candidates);
+      arma::fvec result = swapTarget(
+          data,
+          distMat,
+          medoidIndices,
+          &targets,
+          &bestDistances,
+          &secondBestDistances,
+          assignments,
+          0);
+      estimates.elem(targets) =
+          ((numSamples.elem(targets) % estimates.elem(targets)) +
+           (result * batchSize)) /
+          (batchSize + numSamples.elem(targets));
+      numSamples.elem(targets) += batchSize;
+      arma::fvec adjust(targets.n_rows);
+      adjust.fill(p);
+      // Assume swapConfidence is given in logspace
+      adjust = swapConfidence + arma::log(adjust);
+      arma::fvec confBoundDelta = sigma.elem(targets) %
+                    arma::sqrt(adjust / numSamples.elem(targets));
+      ucbs.elem(targets) = estimates.elem(targets) + confBoundDelta;
+      lcbs.elem(targets) = estimates.elem(targets) - confBoundDelta;
+      candidates = (lcbs < ucbs.min()) && (exactMask == 0);
+    }
+
+    // Perform the medoid switch
+    arma::uword newMedoid = lcbs.index_min();
+    // extract old and new medoids of swap
+    size_t k = newMedoid % medoids->n_cols;
+    size_t n = newMedoid / medoids->n_cols;
+    swapPerformed = (*medoidIndices)(k) != n;
+    steps++;
+
+    if (swapPerformed) {
+      (*medoidIndices)(k) = n;
+      medoids->col(k) = data.col((*medoidIndices)(k));
+    }
+    calcBestDistancesSwap(
+      data,
+      distMat,
+      medoidIndices,
+      &bestDistances,
+      &secondBestDistances,
+      assignments,
+      swapPerformed);
+  }
+}
+}  // namespace km
diff --git a/src/algorithms/fastpam1.cpp b/src/algorithms/fastpam1.cpp
index 0bf765b7..48a78975 100644
--- a/src/algorithms/fastpam1.cpp
+++ b/src/algorithms/fastpam1.cpp
@@ -16,11 +16,13 @@
 #include <unordered_map>
 
 namespace km {
-void FastPAM1::fitFastPAM1(const arma::fmat& inputData) {
+void FastPAM1::fitFastPAM1(
+  const arma::fmat& inputData,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat) {
   data = inputData;
   data = arma::trans(data);
   arma::urowvec medoidIndices(nMedoids);
-  FastPAM1::buildFastPAM1(data, &medoidIndices);
+  FastPAM1::buildFastPAM1(data, distMat, &medoidIndices);
   steps = 0;
   medoidIndicesBuild = medoidIndices;
   arma::urowvec assignments(data.n_cols);
@@ -28,7 +30,7 @@ void FastPAM1::fitFastPAM1(const arma::fmat& inputData) {
   bool medoidChange = true;
   while (iter < maxIter && medoidChange) {
     auto previous{medoidIndices};
-    FastPAM1::swapFastPAM1(data, &medoidIndices, &assignments);
+    FastPAM1::swapFastPAM1(data, distMat, &medoidIndices, &assignments);
     medoidChange = arma::any(medoidIndices != previous);
     iter++;
   }
@@ -39,6 +41,7 @@ void FastPAM1::fitFastPAM1(const arma::fmat& inputData) {
 
 void FastPAM1::buildFastPAM1(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   arma::urowvec* medoidIndices
 ) {
   size_t N = data.n_cols;
@@ -89,11 +92,11 @@ void FastPAM1::buildFastPAM1(
 
 void FastPAM1::swapFastPAM1(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   arma::urowvec* medoidIndices,
   arma::urowvec* assignments
 ) {
   float bestChange = 0;
-  float minDistance = std::numeric_limits<float>::infinity();
   size_t swapIn = 0;
   size_t medoidToSwap = 0;
   size_t N = data.n_cols;
@@ -107,6 +110,7 @@ void FastPAM1::swapFastPAM1(
   // calculate quantities needed for swap, bestDistances and sigma
   KMedoids::calcBestDistancesSwap(
     data,
+    distMat,
     medoidIndices,
     &bestDistances,
     &secondBestDistances,
@@ -124,7 +128,7 @@ void FastPAM1::swapFastPAM1(
 
       // Consider making point i a medoid.
       // The total loss then contains at least one term, -di,
-      // because the loss contribution for point i is reduced to 0
+      // because the loss contribution for point i is reduced from di to 0
       deltaTD.fill(-di);
       // TODO(@motiwari): pragma omp parallel for?
       for (size_t j = 0; j < data.n_cols; j++) {
@@ -165,17 +169,16 @@ void FastPAM1::swapFastPAM1(
 
     // Update the loss and perform the swap if the loss would be improved
     if (bestChange < 0) {
-      minDistance = arma::sum(bestDistances) + bestChange;
       (*medoidIndices)(medoidToSwap) = swapIn;
       calcBestDistancesSwap(
         data,
+        distMat,
         medoidIndices,
         &bestDistances,
         &secondBestDistances,
         assignments,
         swapPerformed);
     } else {
-      minDistance = arma::sum(bestDistances);
       swapPerformed = false;
     }
   }
diff --git a/src/algorithms/kmedoids_algorithm.cpp b/src/algorithms/kmedoids_algorithm.cpp
index 3e69b962..d80bd4c6 100644
--- a/src/algorithms/kmedoids_algorithm.cpp
+++ b/src/algorithms/kmedoids_algorithm.cpp
@@ -14,23 +14,34 @@
 #include "fastpam1.hpp"
 #include "pam.hpp"
 #include "banditpam.hpp"
+#include "banditpam_orig.hpp"
 
 namespace km {
+// NOTE: The order of arguments in this constructor must match that of the
+// arguments in kmedoids_pywrapper.cpp, otherwise undefined behavior can
+// result (variables being initialized with others' values)
 KMedoids::KMedoids(
   size_t nMedoids,
   const std::string& algorithm,
   size_t maxIter,
   size_t buildConfidence,
   size_t swapConfidence,
+  bool useCache,
+  bool usePerm,
+  size_t cacheWidth,
+  bool parallelize,
   size_t seed):
     nMedoids(nMedoids),
     algorithm(algorithm),
     maxIter(maxIter),
     buildConfidence(buildConfidence),
     swapConfidence(swapConfidence),
+    useCache(useCache),
+    usePerm(usePerm),
+    cacheWidth(cacheWidth),
+    parallelize(parallelize),
     seed(seed) {
   KMedoids::checkAlgorithm(algorithm);
-
   // Though we initialize seed from the given parameter,
   // we need to call setSeed to pass it to arma
   KMedoids::setSeed(seed);
@@ -38,20 +49,45 @@ KMedoids::KMedoids(
 
 KMedoids::~KMedoids() {}
 
-void KMedoids::fit(const arma::fmat& inputData, const std::string& loss) {
-  batchSize = fmin(inputData.n_rows, batchSize);
+void KMedoids::fit(
+  const arma::fmat& inputData,
+  const std::string& loss,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat) {
+  numMiscDistanceComputations = 0;
+  numBuildDistanceComputations = 0;
+  numSwapDistanceComputations = 0;
+  numCacheWrites = 0;
+  numCacheHits = 0;
+  numCacheMisses = 0;
+
+  if (distMat) {  // User has provided a distance matrix
+    if (distMat.value().get().n_cols != distMat.value().get().n_rows) {
+      // TODO(@motiwari): Change this to an assertion that is properly raised
+      throw std::invalid_argument("Malformed distance matrix provided");
+    }
+    useDistMat = true;
+  } else {
+    // In case the user is running a new problem without a distance matrix
+    // after running a distance matrix problem
+    useDistMat = false;
+  }
 
   if (inputData.n_rows == 0) {
+    // TODO(@motiwari): Change this to an assertion that is properly raised
     throw std::invalid_argument("Dataset is empty");
   }
+  batchSize = fmin(inputData.n_rows, batchSize);
+
   try {
     KMedoids::setLossFn(loss);
     if (algorithm == "PAM") {
-      static_cast<PAM*>(this)->fitPAM(inputData);
+      static_cast<PAM*>(this)->fitPAM(inputData, distMat);
     } else if (algorithm == "BanditPAM") {
-      static_cast<BanditPAM*>(this)->fitBanditPAM(inputData);
+      static_cast<BanditPAM*>(this)->fitBanditPAM(inputData, distMat);
+    } else if (algorithm == "BanditPAM_orig") {
+      static_cast<BanditPAM_orig*>(this)->fitBanditPAM_orig(inputData, distMat);
     } else if (algorithm == "FastPAM1") {
-      static_cast<FastPAM1*>(this)->fitFastPAM1(inputData);
+      static_cast<FastPAM1*>(this)->fitFastPAM1(inputData, distMat);
     }
   } catch (std::invalid_argument& e) {
     std::cout << e.what() << std::endl;
@@ -107,7 +143,7 @@ size_t KMedoids::getBuildConfidence() const {
 }
 
 void KMedoids::setBuildConfidence(size_t newBuildConfidence) {
-  if (algorithm != "BanditPAM") {
+  if (algorithm != "BanditPAM" && algorithm != "BanditPAM_orig") {
     // TODO(@motiwari): Better error type
     throw "Cannot set buildConfidence when not using BanditPAM";
   }
@@ -119,7 +155,7 @@ size_t KMedoids::getSwapConfidence() const {
 }
 
 void KMedoids::setSwapConfidence(size_t newSwapConfidence) {
-  if (algorithm != "BanditPAM") {
+  if (algorithm != "BanditPAM" && algorithm != "BanditPAM_orig") {
     // TODO(@motiwari): Better error type
     throw "Cannot set buildConfidence when not using BanditPAM";
   }
@@ -135,11 +171,91 @@ size_t KMedoids::getSeed() const {
   return seed;
 }
 
+bool KMedoids::getUseCache() const {
+  return useCache;
+}
+
+void KMedoids::setUseCache(bool newUseCache) {
+  // TODO(@motiwari): Throw an error if not using BanditPAM
+  useCache = newUseCache;
+}
+
+bool KMedoids::getUsePerm() const {
+    return usePerm;
+}
+
+void KMedoids::setUsePerm(bool newUsePerm) {
+  // TODO(@motiwari): Throw an error if not using BanditPAM
+  usePerm = newUsePerm;
+}
+
+size_t KMedoids::getCacheWidth() const {
+  // TODO(@motiwari): Throw an error if not using BanditPAM
+  return cacheWidth;
+}
+
+void KMedoids::setCacheWidth(size_t newCacheWidth) {
+  // TODO(@motiwari): Throw an error if not using BanditPAM
+  cacheWidth = newCacheWidth;
+}
+
+bool KMedoids::getParallelize() const {
+  return parallelize;
+}
+
+// TODO(@motiwari): Change this to const bool newParallelize
+void KMedoids::setParallelize(bool newParallelize) {
+  parallelize = newParallelize;
+}
+
+size_t KMedoids::getDistanceComputations(const bool includeMisc) const {
+    if (includeMisc) {
+        return numMiscDistanceComputations +
+          numBuildDistanceComputations +
+          numSwapDistanceComputations;
+    } else {
+        return numBuildDistanceComputations +
+          numSwapDistanceComputations;
+    }
+}
+
+size_t KMedoids::getMiscDistanceComputations() const {
+    return numMiscDistanceComputations;
+}
+
+size_t KMedoids::getBuildDistanceComputations() const {
+  return numBuildDistanceComputations;
+}
+
+size_t KMedoids::getSwapDistanceComputations() const {
+  return numSwapDistanceComputations;
+}
+
+size_t KMedoids::getCacheWrites() const {
+  return numCacheWrites;
+}
+
+size_t KMedoids::getCacheHits() const {
+  return numCacheHits;
+}
+
+size_t KMedoids::getCacheMisses() const {
+  return numCacheMisses;
+}
+
+size_t KMedoids::getTotalSwapTime() const {
+  return totalSwapTime;
+}
+
+float KMedoids::getTimePerSwap() const {
+  return totalSwapTime / steps;
+}
+
 void KMedoids::setLossFn(std::string loss) {
   // TODO(@motiwari): On setting this, clear the cache and the average loss,
   // assignments, medoids, etc.
   std::for_each(loss.begin(), loss.end(), [](char& c){
-    c = ::tolower(c);
+    c = ::tolower(c);  // TODO(@motiwari): Put something before ::
   });
   // TODO(@motiwari): Change this to a switch
   if (std::regex_match(loss, std::regex("l\\d*"))) {
@@ -177,17 +293,20 @@ std::string KMedoids::getLossFn() const {
 
 void KMedoids::calcBestDistancesSwap(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   const arma::urowvec* medoidIndices,
   arma::frowvec* bestDistances,
   arma::frowvec* secondBestDistances,
   arma::urowvec* assignments,
   const bool swapPerformed) {
-  #pragma omp parallel for
+  #pragma omp parallel for if (this->parallelize)
   for (size_t i = 0; i < data.n_cols; i++) {
     float best = std::numeric_limits<float>::infinity();
     float second = std::numeric_limits<float>::infinity();
     for (size_t k = 0; k < medoidIndices->n_cols; k++) {
-      float cost = KMedoids::cachedLoss(data, i, (*medoidIndices)(k));
+      // 0 for MISC
+      float cost =
+          KMedoids::cachedLoss(data, distMat, i, (*medoidIndices)(k), 0);
       if (cost < best) {
         (*assignments)(i) = k;
         second = best;
@@ -208,14 +327,20 @@ void KMedoids::calcBestDistancesSwap(
 
 float KMedoids::calcLoss(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   const arma::urowvec* medoidIndices) {
   float total = 0;
   // TODO(@motiwari): is this parallel loop accumulating properly?
-  #pragma omp parallel for
+  #pragma omp parallel for if (this->parallelize)
   for (size_t i = 0; i < data.n_cols; i++) {
     float cost = std::numeric_limits<float>::infinity();
     for (size_t k = 0; k < nMedoids; k++) {
-      float currCost = KMedoids::cachedLoss(data, i, (*medoidIndices)(k));
+      float currCost = KMedoids::cachedLoss(
+        data,
+        distMat,
+        i,
+        (*medoidIndices)(k),
+        0);  // 0 for Misc
       if (currCost < cost) {
         cost = currCost;
       }
@@ -229,30 +354,54 @@ float KMedoids::calcLoss(
 
 float KMedoids::cachedLoss(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   const size_t i,
   const size_t j,
-  const bool useCache) {
+  const size_t category,
+  const bool useCacheFunctionOverride
+  ) {
+    // TODO(@motiwari): Change category to an enum
+    if (category == 0) {  // MISC
+        numMiscDistanceComputations++;
+    } else if (category == 1) {  // BUILD
+        numBuildDistanceComputations++;
+    } else if (category == 2) {  // SWAP
+        numSwapDistanceComputations++;
+    } else {
+        // TODO(@motiwari): Throw exception
+    }
+
+  if (this->useDistMat) {
+    return distMat.value().get().at(i, j);
+  }
+
   if (!useCache) {
     return (this->*lossFn)(data, i, j);
   }
 
+  // TODO(@motiwari): Should infer n and m from the size of the cache
   size_t n = data.n_cols;
-  size_t m = fmin(n, ceil(log10(data.n_cols) * cacheMultiplier));
+  size_t m = fmin(n, cacheWidth);
 
   // test this is one of the early points in the permutation
   if (reindex.find(j) != reindex.end()) {
     // TODO(@motiwari): Potential race condition with shearing?
     // T1 begins to write to cache and then T2 access in the middle of write?
     if (cache[(m*i) + reindex[j]] == -1) {
+      numCacheWrites++;
       cache[(m*i) + reindex[j]] = (this->*lossFn)(data, i, j);
     }
+    numCacheHits++;
     return cache[m*i + reindex[j]];
   }
+
+  numCacheMisses++;
   return (this->*lossFn)(data, i, j);
 }
 
 void KMedoids::checkAlgorithm(const std::string& algorithm) const {
   if ((algorithm != "BanditPAM") &&
+      (algorithm != "BanditPAM_orig") &&
       (algorithm != "PAM") &&
       (algorithm != "FastPAM1")) {
     // TODO(@motiwari): Better error type
diff --git a/src/algorithms/pam.cpp b/src/algorithms/pam.cpp
index f78b7e3a..f6503655 100644
--- a/src/algorithms/pam.cpp
+++ b/src/algorithms/pam.cpp
@@ -14,10 +14,12 @@
 #include <unordered_map>
 
 namespace km {
-void PAM::fitPAM(const arma::fmat& inputData) {
+void PAM::fitPAM(
+  const arma::fmat& inputData,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat) {
   data = arma::trans(inputData);
   arma::urowvec medoidIndices(nMedoids);
-  PAM::buildPAM(data, &medoidIndices);
+  PAM::buildPAM(data, distMat, &medoidIndices);
   steps = 0;
   medoidIndicesBuild = medoidIndices;
   arma::urowvec assignments(data.n_cols);
@@ -26,7 +28,7 @@ void PAM::fitPAM(const arma::fmat& inputData) {
   bool medoidChange = true;
   while (i < maxIter && medoidChange) {
     auto previous(medoidIndices);
-    PAM::swapPAM(data, &medoidIndices, &assignments);
+    PAM::swapPAM(data, distMat, &medoidIndices, &assignments);
     medoidChange = arma::any(medoidIndices != previous);
     i++;
   }
@@ -37,6 +39,7 @@ void PAM::fitPAM(const arma::fmat& inputData) {
 
 void PAM::buildPAM(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   arma::urowvec* medoidIndices) {
   size_t N = data.n_cols;
   arma::frowvec estimates(N, arma::fill::zeros);
@@ -74,6 +77,7 @@ void PAM::buildPAM(
 
 void PAM::swapPAM(
   const arma::fmat& data,
+  std::optional<std::reference_wrapper<const arma::fmat>> distMat,
   arma::urowvec* medoidIndices,
   arma::urowvec* assignments) {
   float minDistance = std::numeric_limits<float>::infinity();
@@ -85,6 +89,7 @@ void PAM::swapPAM(
 
   KMedoids::calcBestDistancesSwap(
     data,
+    distMat,
     medoidIndices,
     &bestDistances,
     &secondBestDistances,
diff --git a/src/main.cpp b/src/main.cpp
index a65c33a9..62059ab3 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -21,15 +21,22 @@ int main(int argc, char* argv[]) {
     size_t k;
     int opt;
     int prev_ind;
-    size_t maxIter = 1000;
-    size_t buildConfidence = 1000;
-    size_t swapConfidence = 10000;
+    size_t maxIter = 100;
+    size_t buildConfidence = 3;
+    size_t swapConfidence = 4;
     std::string loss = "L2";
     bool f_flag = false;
     bool k_flag = false;
     const int ARGUMENT_ERROR_CODE = 1;
+    bool useCache = true;
+    bool usePerm = true;
+    int seed = 0;
+    int num_data = 0;
+    bool parallelize = false;
 
-    while (prev_ind = optind, (opt = getopt(argc, argv, "f:l:k:v:s:")) != -1) {
+    // TODO(@motiwari): Use a variadic function signature for this instead
+    while (prev_ind = optind,
+        (opt = getopt(argc, argv, "f:l:k:v:s:cpnw:")) != -1) {
         if ( optind == prev_ind + 2 && *optarg == '-' ) {
         opt = ':';
         --optind;
@@ -50,9 +57,24 @@ int main(int argc, char* argv[]) {
             case 'l':
                 loss = optarg;
                 break;
+            case 's':
+                seed = std::stoi(optarg);
+                break;
             case ':':
                 printf("option needs a value\n");
                 return ARGUMENT_ERROR_CODE;
+            case 'c':
+                useCache = true;
+                break;
+            case 'p':
+                usePerm = true;
+                break;
+            case 'n':
+                num_data = std::stoi(optarg);
+                break;
+            case 'w':
+                parallelize = true;
+                break;
             case '?':
                 printf("unknown option: %c\n", optopt);
                 return ARGUMENT_ERROR_CODE;
@@ -77,14 +99,31 @@ int main(int argc, char* argv[]) {
 
     arma::fmat data;
     data.load(input_name);
+
+    if (num_data < 0) {
+        throw std::invalid_argument(
+        "Error: num_data passed was less than 0");
+    } else if (num_data != 0) {
+        data.resize(num_data, data.n_cols);
+    }
+
     km::KMedoids kmed(
       k,
       "BanditPAM",
       maxIter,
       buildConfidence,
-      swapConfidence);
-    kmed.fit(data, loss);
+      swapConfidence,
+      useCache,
+      usePerm,
+      1000,  // Cache Width
+      parallelize,
+      seed);
+    kmed.fit(data, loss, {});
     for (auto medoid : kmed.getMedoidsFinal()) {
       std::cout << medoid << ",";
     }
+    std::cout << "Final loss: " << kmed.getAverageLoss() << "\n";
+    std::cout << "Num Swap Steps: " << kmed.getSteps() << "\n";
+    std::cout << "Total Swap Milliseconds: " << kmed.getTotalSwapTime() << "\n";
+    std::cout << "Average Swap Milliseconds: " << kmed.getTimePerSwap() << "\n";
 }
diff --git a/src/python_bindings/cache_python.cpp b/src/python_bindings/cache_python.cpp
new file mode 100644
index 00000000..582aa0a5
--- /dev/null
+++ b/src/python_bindings/cache_python.cpp
@@ -0,0 +1,88 @@
+/**
+ * @file cache_python.cpp
+ * @date 2021-08-16
+ *
+ * Defines the functions to cache statistics in the Python bindings.
+ */
+
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <carma>
+#include <armadillo>
+
+#include "kmedoids_pywrapper.hpp"
+
+namespace km {
+size_t km::KMedoidsWrapper::getDistanceComputationsPython(
+  const bool includeMisc) {
+  return KMedoids::getDistanceComputations(includeMisc);
+}
+
+size_t km::KMedoidsWrapper::getMiscDistanceComputationsPython() {
+    return KMedoids::getMiscDistanceComputations();
+}
+
+size_t km::KMedoidsWrapper::getBuildDistanceComputationsPython() {
+    return KMedoids::getBuildDistanceComputations();
+}
+
+size_t km::KMedoidsWrapper::getSwapDistanceComputationsPython() {
+    return KMedoids::getSwapDistanceComputations();
+}
+
+size_t km::KMedoidsWrapper::getCacheWritesPython() {
+  return KMedoids::getCacheWrites();
+}
+
+size_t km::KMedoidsWrapper::getCacheHitsPython() {
+  return KMedoids::getCacheHits();
+}
+
+size_t km::KMedoidsWrapper::getCacheMissesPython() {
+  return KMedoids::getCacheMisses();
+}
+
+void distance_computations_python(pybind11::class_<km::KMedoidsWrapper> *cls) {
+  cls->def("getDistanceComputations",
+           &KMedoidsWrapper::getDistanceComputationsPython);
+}
+
+void misc_distance_computations_python(
+    pybind11::class_<km::KMedoidsWrapper> *cls) {
+    cls->def_property_readonly(
+      "misc_distance_computations",
+      &KMedoidsWrapper::getMiscDistanceComputationsPython);
+}
+
+void build_distance_computations_python(
+    pybind11::class_<km::KMedoidsWrapper> *cls) {
+    cls->def_property_readonly(
+      "build_distance_computations",
+      &KMedoidsWrapper::getBuildDistanceComputationsPython);
+}
+
+void swap_distance_computations_python(
+    pybind11::class_<km::KMedoidsWrapper> *cls) {
+    cls->def_property_readonly(
+      "swap_distance_computations",
+      &KMedoidsWrapper::getSwapDistanceComputationsPython);
+}
+
+void cache_writes_python(pybind11::class_<km::KMedoidsWrapper> *cls) {
+  cls->def_property_readonly(
+      "cache_writes",
+      &KMedoidsWrapper::getCacheWritesPython);
+}
+
+void cache_hits_python(pybind11::class_<km::KMedoidsWrapper> *cls) {
+  cls->def_property_readonly(
+      "cache_hits",
+      &KMedoidsWrapper::getCacheHitsPython);
+}
+
+void cache_misses_python(pybind11::class_<km::KMedoidsWrapper> *cls) {
+  cls->def_property_readonly(
+      "cache_misses",
+      &KMedoidsWrapper::getCacheMisses);
+}
+}  // namespace km
diff --git a/src/python_bindings/fit_python.cpp b/src/python_bindings/fit_python.cpp
index 4edf2357..e5666903 100644
--- a/src/python_bindings/fit_python.cpp
+++ b/src/python_bindings/fit_python.cpp
@@ -7,9 +7,11 @@
  */
 
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 #include <pybind11/numpy.h>
 #include <carma>
 #include <armadillo>
+#include <optional>
 
 #include "kmedoids_pywrapper.hpp"
 
@@ -29,6 +31,7 @@ void km::KMedoidsWrapper::fitPython(
   } catch (pybind11::value_error &e) {
     // Throw it again (pybind11 will raise ValueError)
     // TODO(@motiwari): Make this more informative
+    // TODO(@motiwari): Is this throw redundant with the above?
     throw;
   }
   // if k is specified here, then
@@ -36,7 +39,18 @@ void km::KMedoidsWrapper::fitPython(
   if ((kw.size() != 0) && (kw.contains("k"))) {
     KMedoids::setNMedoids(pybind11::cast<int>(kw["k"]));
   }
-  KMedoids::fit(carma::arr_to_mat<float>(inputData), loss);
+
+  // TODO(@motiwari): Add a comment here explaining the massaging, cleanup
+  if ((kw.size() != 0) && (kw.contains("dist_mat"))) {
+    const pybind11::array_t<float>& distMat =
+      pybind11::cast<const pybind11::array_t<float>>(kw["dist_mat"]);
+    const arma::fmat& tmp_array = carma::arr_to_mat<float>(distMat);
+    KMedoids::fit(carma::arr_to_mat<float>(inputData), loss,
+      std::make_optional<std::reference_wrapper<const arma::fmat>>(tmp_array));
+  } else {
+    // TODO(@motiwari): change std::nullopt to nullopt?
+    KMedoids::fit(carma::arr_to_mat<float>(inputData), loss, std::nullopt);
+  }
 }
 
 void fit_python(pybind11::class_<KMedoidsWrapper> *cls) {
diff --git a/src/python_bindings/kmedoids_pywrapper.cpp b/src/python_bindings/kmedoids_pywrapper.cpp
index a239553c..29f0d67b 100644
--- a/src/python_bindings/kmedoids_pywrapper.cpp
+++ b/src/python_bindings/kmedoids_pywrapper.cpp
@@ -7,6 +7,7 @@
  */
 
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 #include <pybind11/numpy.h>
 #include <carma>
 #include <armadillo>
@@ -16,14 +17,37 @@
 
 namespace km {
 PYBIND11_MODULE(banditpam, m) {
+  // Module functions
   m.doc() = "BanditPAM Python library, implemented in C++";
+  m.attr("__version__") = VERSION_INFO;
+  m.def("get_max_threads",  // TODO(@motiwari): change to get_num_threads
+    &omp_get_max_threads, "Returns max number of threads");
+  m.def("set_num_threads",
+    &omp_set_num_threads, "Set the maximum number of threads");
+
+  // Class functions
   pybind11::class_<KMedoidsWrapper> cls(m, "KMedoids");
-  cls.def(pybind11::init<int, std::string, int, int, int>(),
-          pybind11::arg("n_medoids") = NULL,
-          pybind11::arg("algorithm") = "BanditPAM",
-          pybind11::arg("max_iter") = 1000,
-          pybind11::arg("build_confidence") = 1000,
-          pybind11::arg("swap_confidence") = 10000);
+
+  // Constructor
+  // NOTE: The order of these matters! Otherwise you can get warnings about
+  //  variables not being intialized in the same order as their declarations,
+  //  which can lead to undefined behavior (variables being intialized with
+  //  each others' values). The order here much also match that of the
+  //  constructor in kmedoids_algorithm.*pp
+  cls.def(
+    pybind11::init<int, std::string, int, int, int, bool, bool, int, bool>(),
+    pybind11::arg("n_medoids") = 5,
+    pybind11::arg("algorithm") = "BanditPAM",
+    pybind11::arg("max_iter") = 100,
+    pybind11::arg("build_confidence") = 3,
+    pybind11::arg("swap_confidence") = 4,
+    // TODO(@motiwari): Verify these options are re-used correctly on reset
+    pybind11::arg("use_cache") = true,
+    pybind11::arg("use_perm") = true,
+    pybind11::arg("cache_width") = 1000,
+    pybind11::arg("parallelize") = true);
+
+  // Properties
   cls.def_property("n_medoids",
     &KMedoidsWrapper::getNMedoids, &KMedoidsWrapper::setNMedoids);
   cls.def_property("algorithm",
@@ -34,16 +58,38 @@ PYBIND11_MODULE(banditpam, m) {
     &KMedoidsWrapper::getBuildConfidence, &KMedoidsWrapper::setBuildConfidence);
   cls.def_property("swap_confidence",
     &KMedoidsWrapper::getSwapConfidence, &KMedoidsWrapper::setSwapConfidence);
+  cls.def_property("use_cache",
+    &KMedoidsWrapper::getUseCache, &KMedoidsWrapper::setUseCache);
+  cls.def_property("use_perm",
+    &KMedoidsWrapper::getUsePerm, &KMedoidsWrapper::setUsePerm);
+  cls.def_property("cache_width",
+    &KMedoidsWrapper::getCacheWidth, &KMedoidsWrapper::setCacheWidth);
+  cls.def_property("parallelize",
+    &KMedoidsWrapper::getParallelize, &KMedoidsWrapper::setParallelize);
   cls.def_property("loss_function",
     &KMedoidsWrapper::getLossFn, &KMedoidsWrapper::setLossFn);
   cls.def_property("seed",
     &KMedoidsWrapper::getSeed, &KMedoidsWrapper::setSeed);
+
+  // Other functions
   medoids_python(&cls);
   build_medoids_python(&cls);
   labels_python(&cls);
   steps_python(&cls);
   fit_python(&cls);
   loss_python(&cls);
-  m.attr("__version__") = VERSION_INFO;
+
+  // Cache functions
+  distance_computations_python(&cls);
+  misc_distance_computations_python(&cls);
+  build_distance_computations_python(&cls);
+  swap_distance_computations_python(&cls);
+  cache_writes_python(&cls);
+  cache_hits_python(&cls);
+  cache_misses_python(&cls);
+
+  // Swap timing functions
+  time_per_swap_python(&cls);
+  total_swap_time_python(&cls);
 }
 }  // namespace km
diff --git a/src/python_bindings/swap_times_python.cpp b/src/python_bindings/swap_times_python.cpp
new file mode 100644
index 00000000..3901a04b
--- /dev/null
+++ b/src/python_bindings/swap_times_python.cpp
@@ -0,0 +1,35 @@
+/**
+ * @file cache_python.cpp
+ * @date 2021-08-16
+ *
+ * Defines the functions to cache statistics in the Python bindings.
+ */
+
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <carma>
+#include <armadillo>
+
+#include "kmedoids_pywrapper.hpp"
+
+namespace km {
+size_t km::KMedoidsWrapper::getTotalSwapTimePython() {
+  return KMedoids::getTotalSwapTime();
+}
+
+float km::KMedoidsWrapper::getTimePerSwapPython() {
+  return KMedoids::getTimePerSwap();
+}
+
+void total_swap_time_python(pybind11::class_<km::KMedoidsWrapper> *cls) {
+  cls->def_property_readonly(
+    "total_swap_time",
+    &KMedoidsWrapper::getTotalSwapTimePython);
+}
+
+void time_per_swap_python(pybind11::class_<km::KMedoidsWrapper> *cls) {
+  cls->def_property_readonly(
+    "time_per_swap",
+    &KMedoidsWrapper::getTimePerSwapPython);
+}
+}  // namespace km


Test patch:
diff --git a/.github/workflows/run_linux_tests.yml b/.github/workflows/run_linux_tests.yml
index 28fcdc3a..1e7bdb01 100644
--- a/.github/workflows/run_linux_tests.yml
+++ b/.github/workflows/run_linux_tests.yml
@@ -6,7 +6,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.6", "3.7", "3.8", "3.9"]
+        python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"]
     steps:
       - uses: actions/checkout@v2
 
@@ -18,7 +18,7 @@ jobs:
       - name: Install Ubuntu dependencies
         run: |
           sudo apt update
-          sudo apt install -y build-essential checkinstall libreadline-gplv2-dev libncursesw5-dev libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev 
+          sudo apt install -y build-essential checkinstall libncursesw5-dev libssl-dev libsqlite3-dev tk-dev libgdbm-dev libc6-dev libbz2-dev libffi-dev zlib1g-dev 
           sudo apt install -y clang-format cppcheck
 
       - name: Install Python dependencies 
diff --git a/.github/workflows/run_mac_tests.yml b/.github/workflows/run_mac_tests.yml
index 50105d87..d7ab864f 100644
--- a/.github/workflows/run_mac_tests.yml
+++ b/.github/workflows/run_mac_tests.yml
@@ -6,7 +6,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.6", "3.7", "3.8", "3.9"]
+        # Python 3.11 on Mac errors with "Library not loaded: '@rpath/libarmadillo.11.dylib'"
+        python-version: ["3.7", "3.8", "3.9", "3.10"]  #, "3.11"]
     steps:
       - uses: actions/checkout@v2
 
@@ -21,7 +22,7 @@ jobs:
         run: |
           brew install llvm libomp armadillo
 
-      - name: Install Python dependencies 
+      - name: Install Python dependencies
         run: |
           python -m pip install --upgrade pip
           pip install pytest
@@ -53,16 +54,16 @@ jobs:
           pip install --no-use-pep517 --no-build-isolation -vvv -e .
         env:
           # The default compiler on the Github Ubuntu runners is gcc
-          # Would need to make a respective include change for clang 
-          CPLUS_INCLUDE_PATH: /usr/local/include/carma
+          # Would need to make a respective include change for clang
+          CPLUS_INCLUDE_PATH: /usr/local/include/carma:/usr/local/Cellar/libomp/15.0.2/include:/usr/local/Cellar/libomp/15.0.7/include
 
       - name: Downloading data files for tests
         run: |
           mkdir -p data
+          curl -XGET https://motiwari.com/banditpam_data/MNIST_1k.csv > data/MNIST_1k.csv
           curl -XGET https://motiwari.com/banditpam_data/scRNA_1k.csv > data/scRNA_1k.csv
           curl -XGET https://motiwari.com/banditpam_data/scrna_reformat.csv.gz > data/scrna_reformat.csv.gz
           curl -XGET https://motiwari.com/banditpam_data/MNIST_100.csv > data/MNIST_100.csv
-          curl -XGET https://motiwari.com/banditpam_data/MNIST_1k.csv > data/MNIST_1k.csv
           curl -XGET https://motiwari.com/banditpam_data/MNIST_10k.tar.gz > data/MNIST_10k.tar.gz
           tar -xzvf data/MNIST_10k.tar.gz -C data
           curl -XGET https://motiwari.com/banditpam_data/MNIST_70k.tar.gz > data/MNIST_70k.tar.gz
@@ -70,18 +71,8 @@ jobs:
 
       - name: Run smaller suite of test cases
         run : |
-          pytest tests/test_smaller.py 
+          pytest tests/test_smaller.py
 
       # - name: Run larger suite of test cases
       #   run : |
-      #     pytest tests/test_larger.py 
-
-      - name: Verify that the C++ executable compiles and runs
-        run : |
-          git clone https://github.com/ThrunGroup/BanditPAM
-          cd BanditPAM
-          mkdir build
-          cd build
-          sudo cmake ..
-          sudo make
-          src/BanditPAM -f ../data/MNIST_1k.csv -k 5 -l L2
\ No newline at end of file
+      #     pytest tests/test_larger.py
\ No newline at end of file
diff --git a/.github/workflows/run_windows_tests.yml b/.github/workflows/run_windows_tests.yml
index ad1c507a..5d793e2b 100644
--- a/.github/workflows/run_windows_tests.yml
+++ b/.github/workflows/run_windows_tests.yml
@@ -6,7 +6,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.6", "3.7", "3.8", "3.9"]
+        python-version: ["3.7", "3.8", "3.9", "3.10"]
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python ${{ matrix.python-version }}
diff --git a/scripts/dist_mat_test.py b/scripts/dist_mat_test.py
new file mode 100644
index 00000000..3df02643
--- /dev/null
+++ b/scripts/dist_mat_test.py
@@ -0,0 +1,16 @@
+import numpy as np
+import banditpam
+
+X = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+
+dist_mat1 = np.array(
+    [
+        [0, np.sqrt(27), np.sqrt(108)],
+        [np.sqrt(27), 0, np.sqrt(27)],
+        [np.sqrt(108), np.sqrt(27), 0],
+    ]
+)
+kmed = banditpam.KMedoids(n_medoids=2, algorithm="BanditPAM")
+kmed.fit(X, "L2", dist_mat=dist_mat1)
+print("Number of SWAP steps:", kmed.steps)
+print(kmed.medoids)
diff --git a/tests/constants.py b/tests/constants.py
index 9b4b584c..e3b90d7b 100644
--- a/tests/constants.py
+++ b/tests/constants.py
@@ -16,3 +16,5 @@
 
 PROPORTION_PASSING = 0.8
 SCALING_EXPONENT = 1.2
+
+MILLISECONDS_IN_A_SECOND = 1000
diff --git a/tests/test_initialization.py b/tests/test_initialization.py
new file mode 100644
index 00000000..ad48d60b
--- /dev/null
+++ b/tests/test_initialization.py
@@ -0,0 +1,91 @@
+import unittest
+
+from banditpam import KMedoids
+
+
+class InitializationTests(unittest.TestCase):
+    def test_default_initialization(self):
+        kmed = KMedoids()
+        assert kmed.n_medoids == 5, "n_medoids does not match"
+        assert kmed.algorithm == "BanditPAM", "algorithm does not match"
+        assert kmed.max_iter == 100, "max_iter does not match"
+        assert kmed.build_confidence == 3, "build_confidence does not match"
+        assert kmed.swap_confidence == 4, "swap_confidence does not match"
+        assert kmed.use_cache, "use_cache does not match"
+        assert kmed.use_perm, "use_perm does not match"
+        assert kmed.cache_width == 1000, "cache_width does not match"
+        assert kmed.parallelize, "parallelize does not match"
+
+    def test_initialization_1(self):
+        """
+        Test that we can initialize a KMedoids object with given parameters
+        """
+        kmed = KMedoids(
+            n_medoids=5,
+            algorithm="BanditPAM",
+            max_iter=100,
+            build_confidence=9,
+            swap_confidence=11,
+            use_cache=True,
+            use_perm=True,
+            cache_width=10,
+            parallelize=True,
+        )
+        assert kmed.n_medoids == 5, "n_medoids does not match"
+        assert kmed.algorithm == "BanditPAM", "algorithm does not match"
+        assert kmed.max_iter == 100, "max_iter does not match"
+        assert kmed.build_confidence == 9, "build_confidence does not match"
+        assert kmed.swap_confidence == 11, "swap_confidence does not match"
+        assert kmed.use_cache, "use_cache does not match"
+        assert kmed.use_perm, "use_perm does not match"
+        assert kmed.cache_width == 10, "cache_width does not match"
+        assert kmed.parallelize, "parallelize does not match"
+
+    def test_initialization_2(self):
+        kmed = KMedoids(
+            n_medoids=3,
+            algorithm="FastPAM1",
+            max_iter=1,
+            build_confidence=11,
+            swap_confidence=2,
+            use_cache=False,
+            use_perm=True,
+            cache_width=60,
+            parallelize=False,
+        )
+        assert kmed.n_medoids == 3, "n_medoids does not match"
+        assert kmed.algorithm == "FastPAM1", "algorithm does not match"
+        assert kmed.max_iter == 1, "max_iter does not match"
+        assert kmed.build_confidence == 11, "build_confidence does not match"
+        assert kmed.swap_confidence == 2, "swap_confidence does not match"
+        assert not kmed.use_cache, "use_cache does not match"
+        assert kmed.use_perm, "use_perm does not match"
+        assert kmed.cache_width == 60, "cache_width does not match"
+        assert not kmed.parallelize, "parallelize does not match"
+
+    def test_initialization_3(self):
+        kmed = KMedoids(
+            n_medoids=100,
+            algorithm="PAM",
+            max_iter=10000,
+            build_confidence=11000,
+            swap_confidence=2029,
+            use_cache=True,
+            use_perm=False,
+            cache_width=60000,
+            parallelize=True,
+        )
+        assert kmed.n_medoids == 100, "n_medoids does not match"
+        assert kmed.algorithm == "PAM", "algorithm does not match"
+        assert kmed.max_iter == 10000, "max_iter does not match"
+        assert kmed.build_confidence == 11000, "build_confidence does not" \
+                                               " match"
+        assert kmed.swap_confidence == 2029, "swap_confidence does not match"
+        assert kmed.use_cache, "use_cache does not match"
+        assert not kmed.use_perm, "use_perm does not match"
+        assert kmed.cache_width == 60000, "cache_width does not match"
+        assert kmed.parallelize, "parallelize does not match"
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_larger.py b/tests/test_larger.py
index 3e869301..e86de3b5 100644
--- a/tests/test_larger.py
+++ b/tests/test_larger.py
@@ -15,7 +15,7 @@
     MNIST_SIZE_MULTIPLIERS,
     SCRNA_SIZE_MULTIPLIERS,
     PROPORTION_PASSING,
-    SCALING_EXPONENT
+    SCALING_EXPONENT,
 )
 
 
@@ -74,7 +74,7 @@ def test_time_cases_mnist(self):
             # TODO(@motiwari): Timing test will not work
             # need to compute it over number of steps
             self.assertTrue(
-                runtime < (size_multiplier ** SCALING_EXPONENT) * base_runtime
+                runtime < (size_multiplier**SCALING_EXPONENT) * base_runtime
             )
 
     def test_medium_scrna(self):
@@ -90,9 +90,9 @@ def test_medium_scrna(self):
                 data=data,
                 loss="L1",
                 test_build=False,
-                assert_immediately=False
+                assert_immediately=False,
             )
-        self.assertTrue(count >= PROPORTION_PASSING*NUM_MEDIUM_CASES)
+        self.assertTrue(count >= PROPORTION_PASSING * NUM_MEDIUM_CASES)
 
     def test_various_medium_scrna(self):
         """
@@ -106,11 +106,12 @@ def test_various_medium_scrna(self):
             data = self.scrna.sample(n=size)
             count += bpam_agrees_pam(
                 k=SMALL_K_SCHEDULE[i % N_SMALL_K],
-                data=data, loss="L2",
+                data=data,
+                loss="L2",
                 test_build=False,
-                assert_immediately=False
+                assert_immediately=False,
             )
-        self.assertTrue(count >= PROPORTION_PASSING*NUM_MEDIUM_CASES)
+        self.assertTrue(count >= PROPORTION_PASSING * NUM_MEDIUM_CASES)
 
     def test_time_cases_scrna(self):
         """
@@ -131,7 +132,7 @@ def test_time_cases_scrna(self):
             # TODO(@motiwari): Timing test will not work
             # need to compute it over number of steps
             self.assertTrue(
-                runtime < (size_multiplier ** SCALING_EXPONENT) * base_runtime
+                runtime < (size_multiplier**SCALING_EXPONENT) * base_runtime
             )
 
 
diff --git a/tests/test_smaller.py b/tests/test_smaller.py
index 3213864e..7c5d890a 100644
--- a/tests/test_smaller.py
+++ b/tests/test_smaller.py
@@ -64,7 +64,7 @@ def test_small_scrna(self):
                 test_build=True,
                 assert_immediately=False,
                 use_fp=True,
-                )
+            )
 
             # Test agreement with PAM
             count += bpam_agrees_pam(
@@ -74,9 +74,9 @@ def test_small_scrna(self):
                 test_build=True,
                 assert_immediately=False,
                 use_fp=False,
-                )
+            )
         # Occasionally some may fail due to degeneracy in the scRNA dataset
-        self.assertTrue(count >= 2*PROPORTION_PASSING*NUM_SMALL_CASES)
+        self.assertTrue(count >= 2 * PROPORTION_PASSING * NUM_SMALL_CASES)
 
     def test_small_mnist_known_cases(self):
         """
@@ -106,7 +106,7 @@ def test_small_mnist_known_cases(self):
 
         self.assertEqual(
             sorted(kmed_10.build_medoids.tolist()),
-            [16, 24, 32, 49, 70, 82, 87, 90, 94, 99]
+            [16, 24, 32, 49, 70, 82, 87, 90, 94, 99],
         )
         self.assertEqual(
             sorted(kmed_10.medoids.tolist()),
diff --git a/tests/utils.py b/tests/utils.py
index 1fdea6b8..957a8e76 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -39,23 +39,25 @@ def bpam_agrees_pam(
 
     bpam_and_pam_agree = 1 if bpam_final_medoids == pam_final_medoids else 0
     if test_build:
-        bpam_and_pam_agree &= (bpam_build_medoids == pam_build_medoids)
+        bpam_and_pam_agree &= bpam_build_medoids == pam_build_medoids
 
     if assert_immediately:
-        error_message = ''.join(map(str, [
-                                        "BanditPAM and {} disagree!".format(
-                                            alg_name),
-                                        "\nBanditPAM build medoids:",
-                                        bpam_build_medoids,
-                                        "\n{} build medoids:".format(alg_name),
-                                        pam_build_medoids,
-                                        "\nBanditPAM final medoids:",
-                                        bpam_final_medoids,
-                                        "\n{} final medoids:".format(alg_name),
-                                        pam_final_medoids,
-                                        ]
-                                    )
-                                )
+        error_message = "".join(
+            map(
+                str,
+                [
+                    "BanditPAM and {} disagree!".format(alg_name),
+                    "\nBanditPAM build medoids:",
+                    bpam_build_medoids,
+                    "\n{} build medoids:".format(alg_name),
+                    pam_build_medoids,
+                    "\nBanditPAM final medoids:",
+                    bpam_final_medoids,
+                    "\n{} final medoids:".format(alg_name),
+                    pam_final_medoids,
+                ],
+            )
+        )
         assert bpam_and_pam_agree, error_message
 
     return bpam_and_pam_agree