hw6.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Homework 6: NumPy(29 pts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "name: Robert Frazier\n",
    "\n",
    "email: robertcf@umich.edu\n",
    "\n",
    "This homework assignment took me XX hours in total to complete. (Please help us to gauge the difficulty of the assignment.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Collaboration Disclosure\n",
    "\n",
    "In the cell below, please list *everyone* with whom you discussed any of the homework problems, excluding only the GSIs and the course instructor. \n",
    "\n",
    "If you did not discuss the homework with anyone else, write __\"I did not discuss this homework with anyone.\"__\n",
    "\n",
    "Even if you discuss questions with other, the code you submit must be only yours. All work is checked with the [MOSS plagiarism detector](https://theory.stanford.edu/~aiken/moss/)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "editable": false,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Submission Instructions\n",
    "Your homework solutions should be written entirely in this Jupyter notebook file. Once it contains your solutions, you should submit this notebook through Canvas. \n",
    "\n",
    "\n",
    "Before submitting, please make sure to __Cells->Run All__ executes without errors; errors in your code translate directly to point deductions. \n",
    "In general, you don't need to do explicitly raise errors (e.g. with the ```raise``` function) if we don't ask you to in the problem statement.\n",
    "However, even in cases where we ask you to check for errors, your submission should not contain any examples of your functions actually raising those errors.\n",
    "\n",
    "Note that many parts of this homework where you are expected to type in code will have ```NotImplementedError()``` as a placeholder. You need to delete this function and replace it with your own code.\n",
    "\n",
    "## Homework tips \n",
    "\n",
    "1. **Start early!** If you run into trouble installing things or importing packages, it’s\n",
    "best to find those problems well in advance, not the night before your assignment is\n",
    "due when we cannot help you!\n",
    "\n",
    "2. **Make sure you back up your work!** At a minimum, do your work in a Dropbox\n",
    "folder. Better yet, use git, which is well worth your time and effort to learn.\n",
    "\n",
    "3. **Be careful to follow directions!** Remember that Python is case sensitive. If\n",
    "you are ask you to define a function called my_function and you define a function\n",
    "called My_Function, you will not receive full credit. You may want to copy-paste\n",
    "the function names below to make sure that the functions in your notebook match.\n",
    "\n",
    "## Error checking\n",
    "\n",
    "You do not need to do error checking (raising errors, etc.) in your code unless we explicitly ask you to so in a problem.\n",
    "\n",
    "\n",
    "## Nbgrader\n",
    "\n",
    "We will be using `nbgrader` to grade your jupyter notebook. You will notice some `read-only` cells in the assignment that contain `assert` statements. These are tests that your code must pass for your solution to be correct. If any of the tests fail, you will get an python error and not get points for that question. \n",
    "\n",
    "**Note:** The tests shown not are not comprehensive; additional tests will be used at grading time. You are encouraged to read the problem carefully and verify your code covers all possible cases.\n",
    "\n",
    "**Be careful:** If a jupyter notebook cell takes longer than `60s` to run, the autograder will not grade it and you will receive zero credit for that question."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "markdown",
     "checksum": "b7cffcd79e3179467d5e889abef08484",
     "grade": false,
     "grade_id": "cell-e10bd26a8687d9ce",
     "locked": true,
     "schema_version": 3,
     "solution": false,
     "task": false
    }
   },
   "source": [
    "## Question 1:. The Bakersville Bake-off! (3 pts)\n",
    "\n",
    "The houses in the town of Bakersville are arranged on a grid with R rows (numbered 1 through R) and C columns (numbered 1 through C). The house located on row r and on column c is denoted by (r, c). We say two houses are **neighbors** if they have a common side. \n",
    "\n",
    "During the annual Bakersville Bake-off, each household bakes a certain number of cakes to share with their neighbors.\n",
    "The town comptroller keeps track in an $R \\times C$ matrix $A$, where $A_{r, c}$ is the number of cakes which household $(r, c)$ has baked.\n",
    "\n",
    "The stability of the social order in Bakersville depends on each house baking enough cakes so they can give at least one cake to each neighbor and submit one cake to be judged. Thus, we say a house is **stable** if the number of cakes $A_{r, c}$ is strictly greater than its number of neigbors. We say Bakersville is **stable** if every house is stable. \n",
    "\n",
    "The comptroller has asked you to define a function `is_stable` that takes an $R \\times C$ numpy array (i.e., an ndarray) as its only argument and returns a Python Boolean that is True if and only if the argument is a stable array and False otherwise. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def is_stable(array):\n",
    "    if array.ndim == 1:\n",
    "        array = array.reshape(1, -1)\n",
    "    rows, cols = array.shape\n",
    "    \n",
    "    for i in range(rows):\n",
    "        for j in range(cols):\n",
    "            count = 0\n",
    "            if i > 0:\n",
    "                count += 1\n",
    "            if i < rows - 1:\n",
    "                count += 1\n",
    "            if j > 0:\n",
    "                count += 1\n",
    "            if j < cols - 1:\n",
    "                count += 1\n",
    "            if array[i, j] < count:\n",
    "                return False\n",
    "    return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "dc36504724cf563ca1ff1398b0d04217",
     "grade": true,
     "grade_id": "correct_stable",
     "locked": false,
     "points": 3,
     "schema_version": 3,
     "solution": false
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "array1 = np.array([4, 7, 8,\n",
    "                   5, 10, 9,\n",
    "                   3, 5, 7]).reshape(3, 3)\n",
    "array2 = np.array([3, 2, 4,\n",
    "                   1, 1, 2,\n",
    "                   3, 3, 3]).reshape(3, 3)\n",
    "\n",
    "\n",
    "assert is_stable(array1)\n",
    "assert not is_stable(array2)\n",
    "assert not is_stable(np.array([1., 0., 1.]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 2: Down with `for` loops (12 pts)\n",
    "Each of the problems below contains a function that uses `for` loops to perform a certain operation on arrays. Your job is to re-write these function using only Numpy array manipulations and library functions (e.g. `np.xxx()`). Do not use any `for` or `while` loops, iterators, generators, or list comprehensions in your solutions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**2(a)**  Return all of the rows of the 2-D integer matrix `A` where where each entry of the row is distinct (3 pts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n",
      "[[0 1 2]\n",
      " [3 4 5]\n",
      " [6 7 8]]\n"
     ]
    }
   ],
   "source": [
    "def distinct_rows_py(A):\n",
    "    'Return all rows of A that have completely distinct entries.'\n",
    "    return np.array([a for a in A if len(set(a)) == len(a)])\n",
    "A = np.eye(5)\n",
    "print(distinct_rows_py(A))\n",
    "B = np.arange(9).reshape(3, 3)\n",
    "print(distinct_rows_py(B))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "88f9cd4af96bc3ffafa573cce729e9d8",
     "grade": false,
     "grade_id": "distinct_rows_np",
     "locked": false,
     "schema_version": 3,
     "solution": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "def distinct_rows_np(A):\n",
    "    'Return all rows of A that have completely distinct entries.'\n",
    "    return A[np.all(np.diff(np.sort(A, axis=1), axis=1), axis=1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "6a8e0fe7484f122da08f2c6e59160905",
     "grade": true,
     "grade_id": "correct_distinct",
     "locked": false,
     "points": 3,
     "schema_version": 3,
     "solution": false
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "A = np.eye(5)\n",
    "np.testing.assert_allclose(distinct_rows_np(A), np.empty([0, 5]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**2(b)** Given a vector $v$ of length $n$, and an integer $0<k<=n$, return a 2D array as: (3 pts)\n",
    "```\n",
    "[[v[0], ..., v[k-1]],\n",
    " [v[1], ..., v[k]  ],\n",
    " [v[2], ..., v[k+1]],\n",
    " [ .     .       . ],\n",
    " [ .     .       . ],\n",
    " [ .     .       . ],\n",
    " [v[n-k+1, ..., v[n]]\n",
    " ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 2, 3],\n",
       "       [2, 3, 4],\n",
       "       [3, 4, 5]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def sliding_stack_py(v, k):\n",
    "    \"Stack sliding windows of v of length k.\"\n",
    "    rows = []\n",
    "    for i in range(len(v) - k + 1):\n",
    "        rows.append(v[i : (i + k)])\n",
    "    return np.array(rows)\n",
    "sliding_stack_py(np.array([1, 2, 3, 4, 5]), 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "4f7ee12df40e42c270b4ddcb9027a471",
     "grade": false,
     "grade_id": "sliding",
     "locked": false,
     "schema_version": 3,
     "solution": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "def sliding_stack_np(v, k):\n",
    "    \"Stack sliding windows of v of length k.\"\n",
    "    return np.lib.stride_tricks.sliding_window_view(v, k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "fba971b29747902863514405d7397908",
     "grade": true,
     "grade_id": "correct_sliding",
     "locked": false,
     "points": 3,
     "schema_version": 3,
     "solution": false
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "np.testing.assert_allclose(sliding_stack_np(np.array([1, 2, 3, 4, 5]), 3), np.array([[1, 2, 3],[2, 3, 4],[3, 4, 5]]))\n",
    "np.testing.assert_allclose(sliding_stack_np(np.array([1, 2, 3, 4, 5]), 5), np.array([[1, 2, 3, 4, 5]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**2(c)** Given a vector of non-negative integers `v`, with `max(v) = m`, return a vector `c` of length `m + 1` such that `c[i]` is the number of times that the integer `i` appears in `v`.(3 pts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2, 2, 3, 1, 0, 1])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def digit_count_py(v):\n",
    "    m = max(v)\n",
    "    ret = np.zeros(m + 1, int)\n",
    "    for vv in v:\n",
    "        ret[vv] += 1\n",
    "    return ret\n",
    "v = np.array([0, 0, 1, 1, 2, 2, 2, 3, 5])\n",
    "digit_count_py(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "4b482bcce51a6486678bc7a37be3bd9a",
     "grade": false,
     "grade_id": "digit_count_np",
     "locked": false,
     "schema_version": 3,
     "solution": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "def digit_count_np(v):\n",
    "    return np.bincount(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "d50b09a81c5b7a1da1f111d5f7a64af3",
     "grade": true,
     "grade_id": "correct_digit_count_np",
     "locked": false,
     "points": 3,
     "schema_version": 3,
     "solution": false
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "v = np.array([0, 0, 1, 1, 2, 2, 2, 3, 4])\n",
    "np.testing.assert_allclose(digit_count_np(v), np.array([2, 2, 3, 1, 1]))\n",
    "v = np.array([0, 0, 0, 0, 2, 2, 2, 3, 4])\n",
    "np.testing.assert_allclose(digit_count_np(v), np.array([4, 0, 3, 1, 1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**2(d)** Call a square $n\\times n$ matrix $A$ *countersymmetric* if $A_{ij} = A_{n-j,n-i}$ for all $i$ and $j$. An example of such a matrix is:\n",
    "\n",
    "$$\n",
    "\\begin{pmatrix}\n",
    "4 & 3 & 2 & 1 & 0\\\\\n",
    "8 & 7 & 6 & 5 & 1\\\\\n",
    "11 & 10 & 9 & 6 & 2\\\\\n",
    "13 & 12 & 10 & 7 & 3\\\\\n",
    "14 & 13 & 11 & 8 & 4\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "Write a function `is_countersym` that checks this property: (3 pts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def is_countersym_py(A):\n",
    "    \"Returns True if A is countersymmetric\"\n",
    "    n = A.shape[0]\n",
    "    for i in range(n):\n",
    "        for j in range(n):\n",
    "            if A[i, j] != A[n - j - 1, n - i - 1]:\n",
    "                return False\n",
    "    return True\n",
    "cs_matrix = np.array([[ 4,  3,  2,  1,  0], [ 8,  7,  6,  5,  1], [11, 10,  9,  6,  2], [13, 12, 666,  7,  3], [14, 13, 11,  8,  4]])\n",
    "is_countersym_py(cs_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def is_countersym_np_more(A):\n",
    "    \"Returns True if A is countersymmetric\"\n",
    "    if A.ndim == 1:\n",
    "        A = A.reshape(1, -1)\n",
    "    return np.all(np.tril(np.rot90(A, 1)) == np.rot90(np.tril(np.rot90(A, 3)), 2).T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "b628fd77f1e6a791a62de02a35bb176e",
     "grade": false,
     "grade_id": "is_countersym_np",
     "locked": false,
     "schema_version": 3,
     "solution": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def is_countersym_np_other(A):\n",
    "    if A.ndim == 1:\n",
    "        A = A.reshape(1, -1)\n",
    "    return np.all(A == np.rot90(A, 2).T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_countersym_np(A):\n",
    "    if A.ndim == 1:\n",
    "        A = A.reshape(1, -1)\n",
    "    flipped_matrix = A[:, ::-1]\n",
    "    return np.all(np.triu(flipped_matrix) == np.tril(flipped_matrix).T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "81a8453611f4803da8871fd83ac0202f",
     "grade": true,
     "grade_id": "correct_is_countersym_np",
     "locked": false,
     "points": 3,
     "schema_version": 3,
     "solution": false
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "cs_matrix = np.array([[ 4,  3,  2,  1,  0], [ 8,  7,  6,  5,  1], [11, 10,  9,  6,  2], [13, 12, 10,  7,  3], [14, 13, 11,  8,  4]])\n",
    "assert is_countersym_np(cs_matrix), \"The matrix is not counter-symmetric\"\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Question 3: $k$-means clustering (14 pts)\n",
    "\n",
    "$k$-means is a fundamental algorithm for clustering multivariate data. The inputs to the algorithm are:\n",
    "- An $n\\times p$ data matrix $X$ consisting of $n$ observations of a $p$-dimensional feature vector, and\n",
    "- A $k\\times p$ matrix $C$ containing initial guesses for each $k$ cluster centers.\n",
    "\n",
    "The algorithm proceeds by iteratively a) assigning each point to the nearest cluster center, and b) recomputing the cluster centers as the mean of all of the currently assigned points. Here is a partial implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def kmeans(X, C):\n",
    "    \"\"\"\n",
    "    K-means algorithm.\n",
    "\n",
    "    Args:\n",
    "        - X: ndarray, shape (n, p), n observations of p-dimensional feature vector\n",
    "        - C: ndarray, shape (k, p), k initial cluster centers\n",
    "\n",
    "    Returns:\n",
    "        Tuple of length two:\n",
    "        The first entry is integer ndarray, shape (n), cluster assignments for each data point\n",
    "        The second entry is ndarray, shape (k, p), centers of each cluster\n",
    "    \"\"\"\n",
    "    assert X.shape[1] == C.shape[1]  # p should match\n",
    "    old_assignments = None\n",
    "    while True:\n",
    "        assignments = nearest_cluster(X, C)\n",
    "        if np.all(assignments == old_assignments):\n",
    "            # converged\n",
    "            return assignments, C\n",
    "        old_assignments = assignments\n",
    "        C = compute_centroids(X, assignments)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will finish implementing this algorithm by completing the missing functions `nearest_cluster()` and `compute_centroids()` below. Note: as in the Question 1, avoid using loops, iterators, or comprehensions. You can find solutions with only broadcasting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**3(a)** Implement the function `nearest_cluster`. It should take two array arguments, the data points `X` and the cluster centers `C`, and return an integer array giving the index in `C` which is nearest to each point in `X`. (3 pts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "b30b65131930d9c5ea7a84db535eb3ee",
     "grade": false,
     "grade_id": "3a",
     "locked": false,
     "schema_version": 3,
     "solution": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def nearest_cluster(X, C):\n",
    "    \"\"\"\n",
    "    For each point in X, find the nearest point in C.\n",
    "\n",
    "    Args:\n",
    "        X: ndarray, shape (n, p), n points of dimension p.\n",
    "        C: ndarray, shape (k, p), k points of dimension p.\n",
    "\n",
    "    Returns:\n",
    "        Integer array of length n, [j[1], j[2], ..., j[n]], such that |X[i] - C[j[i]]| <= |X[i] - C[ell]| for 1 <= ell <= k.\n",
    "    \"\"\"\n",
    "    return np.argmin(np.linalg.norm(X[:, None] - C, axis=2), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "568984fe9839146e831caf4ba56b1bbc",
     "grade": true,
     "grade_id": "correct_3a",
     "locked": false,
     "points": 3,
     "schema_version": 3,
     "solution": false
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "ans = np.array([0, 0, 0, 0, 1, 1, 1, 2, 2, 2])\n",
    "X = np.array([[-1], [0.5], [1], [1.3], [1.6], [2], [2.3], [2.7], [3], [100]])\n",
    "C = np.array([[1], [2], [3]])\n",
    "np.testing.assert_allclose(nearest_cluster(X, C), ans)\n",
    "\n",
    "X = np.array([[-101, 1], [1,1], [1, 1.01], [1.5, 1.6], [2.5, -3], [2.5, 2.7], [3, 33]])\n",
    "C = np.array([[1, 1], [2,2], [3,3]])\n",
    "ans = np.array([0, 0, 0, 1, 0, 2, 2])\n",
    "np.testing.assert_allclose(nearest_cluster(X, C), ans)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**3(b)** Implement the function `compute_centroids`. It should take two array arguments, the data points `X` and the assignment array `a`, and return an $k \\times p$ array containing the cluster centroids (averages) for each point assigned to cluster $0, \\dots, k-1$. (You may assume that every entry of $a$ is between $0$ and $k-1$, inclusive.) (3 pts) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_centroids(X, a):\n",
    "    new_centroids = np.array([X[a == i].mean(axis=0) for i in np.unique(a)])\n",
    "    return new_centroids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "4af9b0a2548f6ba45484c7d96af8b347",
     "grade": true,
     "grade_id": "correct_3b",
     "locked": false,
     "points": 3,
     "schema_version": 3,
     "solution": false
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "np.random.seed(1)\n",
    "X = np.random.rand(30, 2)\n",
    "a = np.random.randint(0, 5, size=30)\n",
    "assert compute_centroids(X, a).shape == (5, 2)\n",
    "assert np.allclose(X[a == 0].mean(axis=0), compute_centroids(X, a)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**3(c)** The performance of the $k$-means algorithm is known to depend heavily on the starting point (the initial clusters `C` passed in as the second argument.) In some cases, using a \"good\" starting point can dramatically improve the performance of the algorithm.\n",
    "\n",
    "The $k$-means++ algorithm is designed to find such a good starting point. [According to Wikipedia](https://en.wikipedia.org/wiki/K-means%2B%2B), the steps of $k$-means++ are:\n",
    "\n",
    "1. Choose one center uniformly at random among the data points.\n",
    "2. For each data point $x$ not chosen yet, compute $D(x)$, the distance between $x$ and the nearest center that has already been chosen.\n",
    "3. Choose one new data point at random as a new center, using a weighted probability distribution where a point $x$ is chosen with probability proportional to $D(x)^2$.\n",
    "4. Repeat Steps 2 and 3 until $k$ centers have been chosen.\n",
    "\n",
    "Implement this algorithm using the skeleton provided below. As before, your implementation should only use Numpy functions--no additional loops or comprehensions. \n",
    "\n",
    "**Note**: To ensure reproducibility, the parts of the algorithm that rely on ranndomness are provided for you. Your job is to fill in the missing lines necessary to complete the algoritm. (5 pts.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "a130a2d119d73a5703f4b677dc7879f2",
     "grade": false,
     "grade_id": "3c",
     "locked": false,
     "schema_version": 3,
     "solution": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def kmeanspp(X, k, rng):\n",
    "    \"\"\"\n",
    "    k-means++ algorithm.\n",
    "\n",
    "    Args:\n",
    "        - X: ndarray, shape (n, p), as above.\n",
    "        - k, the number of clusters.\n",
    "        - rng: instance of np.random.Generator().\n",
    "\n",
    "    Returns:\n",
    "        ndarray, shape (k, p), cluster centers.\n",
    "    \"\"\"\n",
    "    n, p = X.shape\n",
    "    C = np.zeros((k, p))\n",
    "    # step 1\n",
    "    j = rng.choice(n)\n",
    "    C[0] = X[j]\n",
    "    for i in range(1, k):\n",
    "        # step 2\n",
    "        D = np.linalg.norm(X - C[:i], axis=1)\n",
    "        w = D**2\n",
    "        w /= w.sum()\n",
    "        # step 3\n",
    "        j = rng.choice(n, p=w)\n",
    "        C[i] = X[j]\n",
    "    return C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "b4211c9685b13a320c76ceb3dd6c2bef",
     "grade": true,
     "grade_id": "correct_3c",
     "locked": false,
     "points": 5,
     "schema_version": 3,
     "solution": false
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.random.seed(1)\n",
    "X = np.random.rand(30, 4)\n",
    "k = 2\n",
    "rng = np.random.default_rng(12345)\n",
    "assert kmeanspp(X, k, rng).shape == (2, 4)\n",
    "ans = np.array([[0.87638915, 0.89460666, 0.08504421, 0.03905478], \n",
    "                [0.6634415 , 0.62169572, 0.11474597, 0.94948926]])\n",
    "np.testing.assert_allclose(kmeanspp(X, k, rng), ans)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**3(d)** In order to measure how good a clustering is, we can define the *within-class variance* \n",
    "\n",
    "$$ V(\\mathbf{X}, \\mathbf{a}, \\mathbf{C}) = \\sum_{i=1}^n \\| \\mathbf{x}_i - \\mathbf{c}_{a_i} \\|^2,$$\n",
    "\n",
    "where the $i$-th element of $\\mathbf{a}=\\{a_1,\\dots,a_n\\}$ is the cluster assignment of observation $i$, and $\\mathbf{C}=(\\mathbf{c}_1,\\dots,\\mathbf{c}_k)$ are the centers of each cluster. Thus, $V(\\mathbf{X}, \\mathbf{a}, \\mathbf{C})$ is the sum of the squared distance from each data point to the center of its assigned cluster.\n",
    "\n",
    "Implement this function. (Again, no loops, just use Numpy functions.) (3 pts) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "43a0fd9d18640089785e840d25790657",
     "grade": false,
     "grade_id": "3d",
     "locked": false,
     "schema_version": 3,
     "solution": true
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def V(X, a, C):\n",
    "    return np.sum(np.linalg.norm(X - C[a], axis=1)**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "cell_type": "code",
     "checksum": "436a72067cc770310f5b82dbd986db3e",
     "grade": true,
     "grade_id": "correct_3d",
     "locked": false,
     "points": 3,
     "schema_version": 3,
     "solution": false
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "X = np.arange(9).reshape(3, 3)\n",
    "a = np.arange(3)\n",
    "C = 2 + X\n",
    "np.testing.assert_allclose(V(X, a, C), 36)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "general",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}