アルゴリズムと数学的思考力

厳しい。年始早々厳しさを感じている。自分のプログラミング力にだ。伸び悩んでいる。

端的に言って、数学力のなさが自分のプログラミング能力に制限をかけている。例えばこの問題。

560. Subarray Sum Equals K
入力として与えられる配列 nums のうち、合計が k となる部分配列の個数を数え上げよ。どうも有名な問題らしいが…

まず大前提として、部分配列なので i, j の2重ループで始点・終点を定めて sum(nums[i, j]) = k になるものを数え上げれば必ず答えが得られる。最悪計算量は O(N^3) ただし i < nums.length < 20000 という制約があるので N^3 では遅すぎるから何か考えてくださいというのがスタート地点。

ここで、結果の変わらない累積和を何度も求めているので nums[i, j] = k を求めたい場合、 nums[0, j] - nums[0, i - 1] としてこれまでの計算結果を利用できることに気付けると計算量を大幅に減らせる。ただし、nums には負数が含まれているので同じ累積和が何度も出現する可能性があるので Map で出現頻度をカウントする

public int subarraySum(int[] nums, int k) {
    int j = 0;
    Map<Integer, Integer> seen = new HashMap<>();
    seen.put(0, 1);
    int counter = 0;
    for (int num : nums) {
        j += num;
        int i = j - k;
        if (seen.containsKey(i)) {
            counter += seen.get(i);
        }
        seen.put(j, seen.getOrDefault(j, 0) + 1);
    }
    return counter;
}

これで時間計算量 O(N) 空間計算量 O(N) で通った。 seen.put(0, 1) の部分だけトリッキーで、このアルゴリズムがそもそも前の計算結果を利用するものなので、初めて sum(nums[0, j]) == k になるケースをカバーする。

これってもはや算数の問題で、区間 nums[i, j] = knums[0, j] - nums[0, i - 1] で求まるっていうのは答え見て初めてマツコデラックスが感心する時の声で「ハ〜〜〜〜〜〜!」って唸ることしかできなかった。もう1題。

373. Find K Pairs with Smallest Sums
ソート済み配列 nums1, nums2 と整数 k が与えられる。このとき nums1, nums2 から1つずつ取り出してつくるペア (u,v) を、和が小さい順に k 個取り出せ。

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]] 

これも O(N^2) で全ペアが列挙でき、それを整列するコストが O(N log N) なので最悪計算量 O(N^2) で必ず解ける問題。これには特に nums.length, k の制約が示されていないが、AtCoderなどのコンテストでは間違いなく O(N^2) では制限時間に収まらない(=不正解となる)データが与えられるだろう。したがってもっと高速なアルゴリズムを考えなさいというのが趣旨となる。

ここで nums1, nums2 はソート済みなので、nums1[0], nums2[0] は常に最小のペアであることが保証される。このような最小のインデックスをそれぞれ i, j とした場合、nums1[i], nums2[j] が解答のひとつならば nums1[i + 1], nums2[j]nums1[i], nums2[j + 1] が次の答えの候補になる。言われてみればそうである(孫権顔)

後は候補が複数あった場合に低コストで最小値を得られるデータ構造を考える。そう、二分ヒープ(最小)である。min heap の追加・削除の最悪コストは O(log N) なので、今回要素数 k の候補から常に最小を取り出しつつ k 個集めるまで繰り返すので O(K log K) が模範解答のひとつとなる。本件は競プロ界で著名な 宇宙ツイッタラーX さんに解説いただいて初めて理解できた。スレッドを参照されたい。

実装例も、氏のそれより綺麗には到底書けないので 氏のgist を見ていただきたい。

で、ここからが主題なのだが、これらの問題はプログラミング能力というか算数の問題なのだ。最初の問題はデータ構造と言っても HashMap ぐらいしか使ってないし、2題目も min heap こそ確かにコンピュータ・サイエンス的なデータ構造ではあるものの、これは最小(最大)値を O(log N) で出し入れできる(最小(最大)値の参照だけなら O(1))データ構造とだけ覚えていれば使える類のものだし、そこがこの問題を解く最大のキーというわけではない。現に僕は min heap を使うと模範解答になるというヒントを貰った後も「 nums1[i], nums2[j] が答えならば nums1[i + 1], nums2[j]nums1[i], nums2[j + 1] がその候補になる」という核心部分が出てこなかったし、むしろここがこの問題の最大のキーポイントであるわけで。

ここからは自分でもうまく整理できていないのだが、これは「数学力」なんだろうか。僕が数学が苦手なのは自他ともに認めるところだが、いま大学院の授業についていくためにマセマの微分積分とかを解いてるけど、なんというかこう、公式あてはめて微分するとかとは違うわけじゃないですか。使われているのは中学生レベルの四則演算で、むしろ「数学的思考力」を問われているというか。

社会人になってプログラマとして10年ぐらいやっていて忘れかけていたけど、そういえば僕は勉強まったくできなかったな、受験敗者だったわ、忘れてた。僕には数学的思考力が致命的にない。こういうのをアルゴリズムクイズを解いていると嫌というほど思い出させられる。

今回の2題は非常に悔しい思いをしたので覚えたが、これを繰り返していけば僕のように頭がよくない人間も類題を解けるようになるのか、それともちょっとひねられるだけでまたまったくお手上げになってしまって悔し涙に枕を濡らすのか、いまは分からない。願わくは2年後に読み返したときに「あの時はあんなしょうもない問題で悩んでいたんだなぁ😅」と思えるようになっていたいが、どうやることやら。正月に書くような内容ではないな…