Pythonでターゲット配列を作成するための最小関数呼び出し回数を求めるプログラム

問題の概要

次のようなPythonの関数があるとします。

def modify(arr, op, index):
    if op == 0:
        arr[index] += 1
    if op == 1:
        for i in range(len(arr)):
            arr[i] *= 2

この関数は、op=0 のとき指定されたインデックスの要素を1増やし、op=1 のとき配列内のすべての要素を2倍にします。

同じサイズのゼロ配列 [0, 0, ..., 0] を出発点として、目的の配列 nums を作り上げるまでに必要な関数呼び出しの最小回数を求めるのがこの問題です。

入力例と動作の流れ

たとえば、入力が nums = [1, 5, 3] の場合、答えは 7 になります。初期状態はすべての要素が0の [0, 0, 0] です。

• 2番目の要素を1増やす → [0, 1, 0](1回)
• 2番目の要素を2倍にする → [0, 2, 0](1回)
• 3番目の要素を1増やす → [0, 2, 1](1回)
• インデックス1以降の要素をまとめて2倍にする → [0, 4, 2](1回)
• 残りの要素をそれぞれ1ずつ増やす → [1, 5, 3](3回)

合計 4 + 3 = 7回 の操作が必要であることがわかります。

解き方の考え方

この問題を効率的に解くには、各要素を「目標値から逆算」していくアプローチが有効です。ポイントは次の2つです。

• 加算操作(op=0)は特定の1要素にしか影響しないため、各要素ごとに必要な回数を数えて合計する
• 倍加操作(op=1)は配列全体に一度に作用するため、全要素の中で最も多く倍加が必要な要素の回数だけ数えればよい

具体的には、各数値 n に対して以下の処理を繰り返します。

• n が偶数なら、n を2で割り、倍加カウント(double)を1増やす
• n が奇数なら、n から1を引き、加算カウント(ans[0])を1増やす

これを n が0になるまで続けると、加算回数は「n を2進数で表したときの1の個数」に一致し、倍加回数は「ビット長から1を引いた値」に一致します。最終的な答えは、加算回数の合計 + 全体で必要な最大倍加回数となります。

実装例(Python)

def solve(nums):
    ans = [0, 0]
    for n in nums:
        double = 0
        while(n):
            if not n % 2:
                n = n // 2
                double += 1
            else:
                n -= 1
                ans[0] += 1
            ans[1] = max(ans[1], double)
    return sum(ans)

nums = [1, 5, 3]
print(solve(nums))

入力

[1, 5, 3]

出力

7

計算量

各要素の処理は、値を2で割り続ける形になるため O(log n) 時間で完了します。したがって、配列全体の時間計算量は O(N log M)(N は要素数、M は最大値)、空間計算量は O(1) です。