Codility Lesson17-1 MinAbsSumをO(N * M^2)で解く

公開:2021/02/11
更新:2021/02/12
3 min読了の目安(約3100字TECH技術記事

Codility内でChallenge以外だとおそらく最難関な気がする。Ambitiousカテゴリがこれ以外だと「PolygonConcavityIndex」だけな為たぶんそう。

問題↓

https://app.codility.com/programmers/lessons/17-dynamic_programming/min_abs_sum/

何が難しいかというと、O(N^2 * M)ならば割と誰でも辿り着けそうなのだが、それではTimeoutになる。(昔はこんなテストケースきつくなかったらしくて、解答後に世界中の記事をみると結構O(N^2 * M)の記事が残ってた)。日本語解説記事が0だったので記事にした。

基本的な方針

  1. 簡単に書けば「Aの要素に1or-1をかけて、総和が一番小さい正の数を作れ」
  2. 1or-1の選択は自由なので、全てのAの要素A[i]は正にも負にもできる
  3. [P:正にしたA[p]の総和] + [Q:負にしたA[q]の総和]が0に近ければいい。
  4. abs(P)とabs(Q)の和は常にsum(A)なので、Pはsum(A)/2に近ければ良い
  5. sum(A)/2に最も近くなるPの組み合わせを探そう
  6. ほぼ部分和問題なので動的計画法を使う

駄目な解法:O(N^2 * M)

いっちばん最初に愚直にdpしたら思いつく解法。愚直の中でもNを最適化しよう四苦八苦したバージョン。

# A = [1, 5, 2, -2]としてコメント残します
def solution(A):
    if len(A) == 0:
        return 0
    A = [abs(i) for i in A] # 1, 5, 2, 2
    S = sum(A) # 10
    target = S//2 # 5

    # その部分和を作れる場合は0以上、作れない場合は-1になるdp
    dp = [-1] * (target + 1)
    dp[0] = 0
    for candidate in A:
        for cur_target in range(target,-1,-1):
            if dp[cur_target] != -1 and cur_target + candidate < target:
                dp[cur_target + candidate] = 1
    '''
    今回の場合、
    Aが[1, 5, 2, 2]だと 
    dpは[0, 1, 1, 1, 1, 1]となる
    '''
    result = S # こいつにはならないだろうって数字。float('inf')でもよし
    for i in range(target, -1, -1):
        # targetに近い方が正解なのでtargetから始める
        if dp[i] >= 0: # その和の組み合わせが可能な場合
            result = min(result, S - 2 * i) 
            # Pの和がS-i, Qの和がS - (S-i)
            # その為P + (-Q) = S - 2 * i となる
            if result == 0: # 0ならそれ未満はありえないので終わり
                break;
    return result

で、これではおそすぎるのである。Aの制約は-100<=A<=100なのでかなりの重複があるのがわかる。dpを工夫する。

正解:O(N * M^2)

def solution(A):
    if len(A) == 0:
        return 0
    A = [abs(i) for i in A]
    S = sum(A)
    M = max(A)
    target = S//2

    # 各数字の重複回数をhashでまとめる
    item_occurancy = {}
    for i in A:
        item_occurancy[i] = item_occurancy.get(i, 0) + 1

    dp = [-1] * (target + 1) 
    dp[0] = 0
    for candidate in item_occurancy: # 重複を無視します
        for cur_target in range(target + 1):  # 重複を無視するので先程と違い0から。
            if dp[cur_target] >= 0:
	        '''
                最初の1手。まず最初にdp[0]にしか置けない。
		ここにその要素の出現回数を置く。
		先程は重複分毎回loopしてたがここで一括で行える
		'''
                dp[cur_target] = item_occurancy[candidate]
            elif (cur_target >= candidate
                  and dp[cur_target - candidate] > 0):
                '''
                仮に1が1つ、3が2つある場合、
                0 2 -1 -1 -1 -1 -1 -1
                           |        |
                           ここと   ここが条件分岐の発生位置になる
		出現回数をそのまま1loopで付与して可能性のある総和の位置にフラグを立てられる。
                0 0 -1 -1  0 -1 -1  0
                '''
                dp[cur_target] = dp[cur_target - candidate] - 1

    # あとはさっきと同じ
    result = S
    for i in range(target, -1, -1):
        if dp[i] >= 0:
            result = min(result, S - 2 * i) 
            if result == 0:
                break;
    return result

クリア↓

まとめ

やり終えたら、そして言語化できたら「そらこうだよ」という気持ちなんだけど、やってる時は最初の愚直dpでtimeoutした時点で「え、dp自体を最適化しないといけないの・・・?あぁtarget以上は見なくていいのか・・・え、まだtimeoutしちゃうの・・・え・・・?」って迷走してショートしました。ハマった時冷静になれないのはまだ脳のHP不足&経験不足なので精進。