Dynamic Programming/SubarraySum_DivByK.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SubarraySum_DivByK.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "a3hYgaHPUkua"
      },
      "source": [
        "'''\n",
        "  Find the number of continuous subarrays of an array whose sum is divisible by k. \n",
        "  Time complexity = O(n)\n",
        "\n",
        "  Parameters:\n",
        "  -----------\n",
        "    nums: list\n",
        "          The input array\n",
        "    k   : int\n",
        "          Given divisor of subarrays\n",
        "\n",
        "  Returns:\n",
        "  --------\n",
        "    count: int\n",
        "           The number of subarrays\n",
        "\n",
        "  Examples:\n",
        "  --------- \n",
        "    Given an array [4,5,0,-2,-3,1] with given divisor k = 5. \n",
        "    Subarrays that have sum % 5 = 0 are:\n",
        "    [5], [5,0], [0], [5,0,-2,-3], [0,-2,-3], [-2,-3], [4,5,0,-2,-3,1].\n",
        "\n",
        "    How the algorithm works?\n",
        "\n",
        "    The array of cumulative sum is [0,4,9,9,7,4,5]\n",
        "    The array of cumulative sum mod K is [0,4,4,4,2,4,0]\n",
        "    In the algorithm, turn the array into dictionary of cumulative sum mod K and its frequency:\n",
        "      memo = {0: 2, 4: 4, 2: 1}\n",
        "\n",
        "    For every cumulative sum (called current sum), if (currSum % k) is in memo, \n",
        "    then we found subarray(s) having subarraySum % k = 0 from index of \n",
        "    the sum = currSum - subarraySum to index of currSum.\n",
        "\n",
        "                      currSum % k = (currSum - subarraySum) % k\n",
        "                              => subarraySum % k = 0\n",
        "\n",
        "    The number of subarray(s) found is the frequency of the sum = currSum % k.\n",
        "\n",
        "    i = 4  -> currSum = 4 -> currSum % k = 4 not in memo -> ignore\n",
        "    i = 5  -> currSum = 9 -> currSum % k = 4 in memo     -> count = 1, found [5]\n",
        "    i = 0  -> currSum = 9 -> currSum % k = 4 in memo     -> count = 3, found [5,0], [0]\n",
        "    i = -2 -> currSum = 7 -> currSum % k = 2 not in memo -> ignore\n",
        "    i = -3 -> currSum = 4 -> currSum % k = 4 in memo     -> count = 6, found [5,0,-2,-3], [0,-2,-3], [-2,-3]\n",
        "    i = 1  -> currSum = 5 -> currSum % k = 0 in memo     -> count = 7, found [4,5,0,-2,-3,1].\n",
        "\n",
        "    >>> nums = [4,5,0,-2,-3,1]\n",
        "    >>> k = 5\n",
        "    >>> print(SubarraySum_DivByK(nums, k))\n",
        "    7\n",
        "\n",
        "  References:\n",
        "    https://leetcode.com/problems/subarray-sums-divisible-by-k/solution/\n",
        "'''\n",
        "\n",
        "def SubarraySum_DivByK(nums, k):\n",
        "  from collections import defaultdict\n",
        "    \n",
        "  memo = defaultdict(int) # Memoization as: cumulativeSum mod K -> Frequency\n",
        "  memo[0] = 1\n",
        "    \n",
        "  currSum = 0  # Current cumulative sum \n",
        "  count = 0 # Number of continuous subarrays whose sum is divisible by K\n",
        "    \n",
        "  for i in nums:\n",
        "    currSum += i     \n",
        "      \n",
        "    # If (currSum % k) is in memo, then we found subarray(s) having subarraySum % k = 0\n",
        "    #   from index of the sum = currSum - subarraySum to index of currSum.\n",
        "    # currSum % k = (currSum - subarraySum) % k\n",
        "    #   => subarraySum % k = 0\n",
        "    if currSum % k in memo:\n",
        "      # The number of subarray(s) found is the frequency of the sum = currSum % k.\n",
        "      count += memo[currSum % k]\n",
        "      \n",
        "    memo[currSum % k] += 1\n",
        "  \n",
        "  return count"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}