Create subarray-sums-divisible-by-k.py

Author: kamyu104

Python/subarray-sums-divisible-by-k.py

@@ -0,0 +1,21 @@
+# Time:  O(n)
+# Space: O(k)
+
+import collections
+
+
+class Solution(object):
+    def subarraysDivByK(self, A, K):
+        """
+        :type A: List[int]
+        :type K: int
+        :rtype: int
+        """
+        count = collections.defaultdict(int)
+        count[0] = 1
+        result, prefix = 0, 0
+        for a in A:
+            prefix = (prefix+a) % K
+            result += count[prefix]
+            count[prefix] += 1
+        return result