Pythonで数値ペアの集合から式の最大値を求めるアルゴリズムと実装例

問題の概要

同じ要素数 N を持つ2つの配列 nums1 と nums2 が与えられているとします。ここで、1 から N までの整数からなる集合 S を考えます。S の空でない部分集合 {i1, i2, ..., ik} を選んだとき、次の式の値を最大化するのがこの問題の目的です。

(nums1[i1] + nums1[i2] + ... + nums1[ik])² + (nums2[i1] + nums2[i2] + ... + nums2[ik])²

具体例

たとえば、入力が nums1 = [-1, 6]、nums2 = [5, 4] の場合、出力は 106 になります。これは次の3通りの選び方を比較した結果です。

• (-1)² + 5² = 26
• 6² + 4² = 52
• (-1 + 6)² + (5 + 4)² = 25 + 81 = 106

このケースでは、すべての要素を選んだ場合が最も大きな値となり、これが答えになります。

解法のアプローチ:幾何学的な捉え方

この問題を効率的に解く鍵は、各ペア (nums1[i], nums2[i]) を2次元平面上のベクトルとみなすことです。部分集合に対する式の値は、「選んだベクトルをすべて足し合わせた合成ベクトルの長さ(ノルム)の2乗」に他なりません。つまり、この問題は「与えられたベクトルの中からいくつかを選び、合成ベクトルの長さを最大化する」という幾何学の問題へと置き換えられます。

さらに重要な性質として、最適な解を構成するベクトル群は、偏角(原点から見た角度)でソートしたときに連続した区間を形成することが知られています。この性質を利用すれば、すべての部分集合を列挙する O(2^N) の総当たりではなく、ソートと貪欲法を組み合わせた O(N²) のアルゴリズムで答えを求められます。

アルゴリズムの手順

具体的には、次の手順で解きます。

1. vs を、各 i についてのペア (nums1[i], nums2[i]) のリストとして作成します。
2. vs を、atan2(v[1], v[0]) で計算される偏角の昇順にソートします。
3. best を 0 で初期化します。
4. i を 0 から vs のサイズ - 1 まで動かしながら、以下を繰り返します。
   • u を vs[i]、l を u[0]² + u[1]² で初期化します。
   • vs を2回連結したリストのうち、インデックス i+1 から i + vsのサイズ - 1 までの各 v について、t1 = (u[0]+v[0], u[1]+v[1])、t2 = t1[0]² + t1[1]² を計算し、t2 ≥ l であれば u と l を更新します。
   • ループ終了後、l > best であれば best を更新します。
   • 同じ処理を、今度は連結リストを逆順にたどって再度行い、best を更新します。
5. 最後に best を返します。

各開始点から順方向・逆方向の両方で貪欲にベクトルを加算していくことで、角度的に連続する最適な区間を取りこぼさずに見つけられるようにしています。

実装例(Python)

理解を深めるために、以下の実装を見てみましょう。

from math import atan2

def solve(nums1, nums2):
    vs = list(zip(nums1, nums2))
    vs = sorted(vs, key=lambda v: atan2(v[1], v[0]))

    best = 0
    for i in range(len(vs)):
        # 順方向(時計回り)への貪欲探索
        u = vs[i]
        l = u[0]*u[0] + u[1]*u[1]
        for v in (vs + vs)[i+1:(i+len(vs))]:
            t1 = (u[0]+v[0], u[1]+v[1])
            t2 = t1[0]*t1[0] + t1[1]*t1[1]
            if t2 >= l:
                u = t1
                l = t2
        if l > best:
            best = l

        # 逆方向(反時計回り)への貪欲探索
        u = vs[i]
        l = u[0]*u[0] + u[1]*u[1]
        for v in reversed((vs + vs)[i+1:(i+len(vs))]):
            t1 = (u[0]+v[0], u[1]+v[1])
            t2 = t1[0]*t1[0] + t1[1]*t1[1]
            if t2 >= l:
                u = t1
                l = t2
                if l > best:
                    best = l
    return best

nums1 = [-1, 6]
nums2 = [5, -4]
print(solve(nums1, nums2))

入力

[-1, 6], [5, -4]

出力

52

出力の検証

このケースでは、考えられる選び方ごとに値を比較すると次のようになります。

• (-1)² + 5² = 26
• 6² + (-4)² = 36 + 16 = 52
• (-1 + 6)² + (5 + (-4))² = 25 + 1 = 26

よって最大値は 52 となり、プログラムの出力と一致します。

まとめ

本記事では、2つの配列から部分集合を選び、(Σnums1)² + (Σnums2)² の最大値を求める問題を扱いました。ポイントは次のとおりです。

• 各ペアを2次元ベクトルとみなし、問題を「合成ベクトルのノルムの最大化」に帰着させる。
• 偏角(atan2)でソートすると、最適解は角度的に連続した区間になるという性質を利用する。
• 配列を2周分に拡張し、各開始点から順方向・逆方向に貪欲にベクトルを加算しながら最大値を更新する。
• 計算量はソートに O(N log N)、探索に O(N²) であり、全列挙の O(2^N) に比べて大幅に高速。