560 subarray sum equals k

560-Subarray-Sum-Equals-K/note.md

@@ -0,0 +1,50 @@
+# 560. Subarray Sum Equals K
+
+<https://leetcode.com/problems/subarray-sum-equals-k/>
+
+## step1（まず通す）
+
+### サブアレイの数はいくつあるか
+
+開始位置と終了位置を選べばサブアレイが決まるので、配列長さ n に対して
+`nC2 = n(n-1) /2`
+サブアレイを全部見るだけで計算量がすごいことになりそう
+
+よくあるのは、先頭から一つずつ伸ばして和を計算（累積和）する方法。これなら O(n)。
+先頭から始まるサブアレイ限定で {和: サブアレイ数} を持つと和の部分の組み合わせを考えて条件を満たすサブアレイ数が計算できそう。
+ただし、サブアレイの終端位置が小さいものから大きいものは引けないので注意する必要がある。
+
+書いてみたけど、計算時間はどれくらいだろう？
+一番気にしてるのは、後半の3重ループが長くなりすぎないかということだが、
+cumulative_sum_to_tail のキー数を x とすると、各キーに対応する value であるリストの平均長さは約 n/x なので、3重ループは平均で `x * n/x * n/x = n^2/x` 回程度？
+各要素の符号が揃っていて累積和が単調だとループ回さず掛け算で済むので楽になりそう。
+
+`nums = [1,-1] * 10 ^ 4, k = 1` みたいな時には 2 * 10 ^ 8 回ループが回るわけで、python だと数秒以上かかりかねないのでは、という気がする。
+
+実際、leetcode のオンラインジャッジで 10 秒くらいかかってる
+
+あと complement は （sum を基準とした時の k の）補数なので命名としてちょっと気持ち悪い。sum の〇〇という命名の方がスッキリくる。
+
+## step2（整形＆他の人のコードを読む）
+
+他の人のコードを読む前に、cumulative_sum_to_tail の values のリストはソートされているのでインデックスの大小関係を確認するのに i,j で2重ループを回す必要がないだろうということに気がついた。
+
+思いつく中では短い方でループを回してバイナリサーチが一番早そう。
+`x * n/x * log(n/x) = n (log n - logx) < n log n` 回くらいにはなる。
+
+n が 2 * 10 ^ 4 なので `n log n ~ 30 * 10 ^ 5 で` 数百 ms くらいにはなるか。
+
+### 他の方のコードを見る
+
+- <https://github.com/shintaro1993/arai60/pull/20/changes>
+  - 累積和を計算するのと一緒に差が k になるものを探せば、そのとき新たに和を計算したものより短い prefix array しかないので長さの比較が不要になり O(n) にできる。
+  - こうすると、prefix array の末尾のインデックスを覚えている必要もなくなるので単にカウントでよくなる
+  - `defaultdict` の初期値は初期化の時に引数で渡せる
+  - `itertools.accumulate()` という関数があるようだ
+    - これに限らず [itertools](https://docs.python.org/ja/3/library/itertools.html)は非常に便利そう
+- <https://github.com/shintaro1993/arai60/pull/20/changes>
+  - ↑と同じようにされている
+  - 自分も for ループの先頭でハッシュマップをインクリメントするよりは、リスト先頭の直前に空の prefix array を想定して初期化する方が好き
+- prefix sum という命名が多め？に感じる
+
+## step3（10分以内にさっとかける * 3回）

560-Subarray-Sum-Equals-K/step1.py

@@ -0,0 +1,24 @@
+from collections import defaultdict
+
+
+class Solution:
+    def subarraySum(self, nums: list[int], k: int) -> int:
+        cumulative_sum_to_tail = defaultdict(list)
+        # 先頭から始まるサブアレイ自体も考慮するために和が0の空アレイを想定
+        cumulative_sum_to_tail[0].append(-1)
+        cumulative_sum = 0
+        for i, num in enumerate(nums):
+            cumulative_sum += num
+            cumulative_sum_to_tail[cumulative_sum].append(i)
+
+        num_sum_equals_k = 0
+        for sum in cumulative_sum_to_tail:
+            complement = sum - k
+            if complement not in cumulative_sum_to_tail:
+                continue
+            for i in cumulative_sum_to_tail[sum]:
+                for j in cumulative_sum_to_tail[complement]:
+                    if i > j:
+                        num_sum_equals_k += 1
+
+        return num_sum_equals_k

560-Subarray-Sum-Equals-K/step2.py

@@ -0,0 +1,13 @@
+from collections import defaultdict
+from itertools import accumulate
+
+
+class Solution:
+    def subarraySum(self, nums: list[int], k: int) -> int:
+        prefix_sum_count = defaultdict(int, {0: 1})
+        equals_k_count = 0
+        for prefix_sum in accumulate(nums):
+            equals_k_count += prefix_sum_count[prefix_sum - k]
+            prefix_sum_count[prefix_sum] += 1
+
+        return equals_k_count

560-Subarray-Sum-Equals-K/step3.py

@@ -0,0 +1,13 @@
+from collections import defaultdict
+from itertools import accumulate
+
+
+class Solution:
+    def subarraySum(self, nums: list[int], k: int) -> int:
+        prefix_sum_count = defaultdict(int, {0: 1})
+        equals_k_count = 0
+        for prefix_sum in accumulate(nums):
+            equals_k_count += prefix_sum_count[prefix_sum - k]
+            prefix_sum_count[prefix_sum] += 1
+
+        return equals_k_count