Friends Code Challenge.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def generate_test_case(counts):\n",
    "    subgraph_counts = np.array(counts)\n",
    "    subgraph_counts_cum = np.cumsum(subgraph_counts)\n",
    "    nodes_total = np.sum(subgraph_counts)\n",
    "    nodes = np.array(range(nodes_total))\n",
    "    np.random.shuffle(nodes)\n",
    "    slices = [slice(j - i, j) for i, j in zip(subgraph_counts, subgraph_counts_cum)]\n",
    "    nodes_sliceds = [nodes[s] for s in slices]\n",
    "    \n",
    "    graph = np.zeros((nodes_total, nodes_total))\n",
    "    \n",
    "    for nodes_sliced in nodes_sliceds:\n",
    "        nodes_sliced_set = set(nodes_sliced)\n",
    "        nodes_chosen = set()\n",
    "        while nodes_chosen != nodes_sliced_set:\n",
    "            nodes_to_link = np.random.choice(nodes_sliced, 2, replace=False)\n",
    "            if (not nodes_chosen) or (nodes_to_link[0] in nodes_chosen) or (nodes_to_link[1] in nodes_chosen):\n",
    "                nodes_chosen.update(nodes_to_link)\n",
    "                graph[nodes_to_link[0], nodes_to_link[1]] = 1\n",
    "                graph[nodes_to_link[1], nodes_to_link[0]] = 1\n",
    "    \n",
    "    graph_text = np.vectorize(lambda x: 'Y' if x == 1 else 'N')(graph)\n",
    "    graph_strings = [''.join(c) for c in graph_text]\n",
    "    return graph_strings\n",
    "\n",
    "test_data_1 = generate_test_case([3, 3, 5])\n",
    "test_data_2 = generate_test_case([2, 3, 5, 3])\n",
    "test_data_3 = generate_test_case(range(2, 50))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "# essentially a sparse matrix\n",
    "class SparseMatrix:\n",
    "    def __init__(self, adjacency_list):\n",
    "        self.adjacency_list = adjacency_list\n",
    "        self.subgraph_count = None\n",
    "    \n",
    "    def count(self):\n",
    "        return len(self.adjacency_list)\n",
    "    \n",
    "    def from_input(input_strings):\n",
    "        adjacency_list = tuple(frozenset([index for index, elem in enumerate(line) if elem == 'Y']) for line in input_strings)\n",
    "        return SparseMatrix(adjacency_list)\n",
    "    \n",
    "    def count_subgraphs(self):\n",
    "        # cached, and this works because frozen sets and immutability fuck yeah\n",
    "        if not self.subgraph_count:\n",
    "            self.subgraph_count = SparseMatrix.__count_subgraphs(self, set(range(len(self.adjacency_list))))\n",
    "        return self.subgraph_count\n",
    "    \n",
    "    def __getitem__(self, node):\n",
    "        if node < self.count():\n",
    "            return self.adjacency_list[node]\n",
    "        return []\n",
    "    \n",
    "    def __count_subgraphs(graph, nodes_unvisited):\n",
    "        subgraph_count = 0\n",
    "        while nodes_unvisited:\n",
    "            subgraph_count += 1\n",
    "            nodes_in_subgraph = set([nodes_unvisited.pop()])\n",
    "            while nodes_in_subgraph:\n",
    "                node = nodes_in_subgraph.pop()\n",
    "                if node in nodes_unvisited:\n",
    "                    nodes_unvisited.remove(node)\n",
    "                neighbors_unvisited = graph[node] & nodes_unvisited\n",
    "                nodes_in_subgraph.update(neighbors_unvisited)\n",
    "        return subgraph_count\n",
    "    \n",
    "def friend_circles(friends_input):\n",
    "    friendships = SparseMatrix.from_input(friends_input)\n",
    "    return friendships.count_subgraphs()\n",
    "\n",
    "assert friend_circles(test_data_1) == 3\n",
    "assert friend_circles(test_data_2) == 4\n",
    "assert friend_circles(test_data_3) == 48"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}