Pythonでスタックのリストからk個の要素をポップしたときの最大合計を求めるプログラム

複数のスタック(リスト)と整数 k が与えられたとき、各スタックから要素をポップし、合計でちょうど k 個の要素を取り出した場合に得られる最大の合計値を求める問題を考えてみましょう。

例えば、入力が stacks = [[50, -4, -15], [2], [6, 7, 8]]、k = 4 の場合、出力は 39 になります。最初のスタックから3つの要素をすべてポップし、最後のスタックから末尾の要素を1つポップすると、-15 + (-4) + 50 + 8 = 39 という合計が得られるためです。

解法のアプローチ

この問題は再帰(バックトラッキング)を用いて解くことができます。重要なポイントは、スタックからは必ず「末尾(積まれている側)」から連続して要素を取り出す必要があるという点です。そのため、あるスタックから j 個取る場合の候補は「末尾 j 個の合計」だけに限定され、探索範囲を絞り込めます。

アルゴリズムの手順

再帰関数 rec(i, n) を定義します。i は現在処理中のスタックのインデックス、n はこれまでにポップした要素数です。

1. n == k の場合:0 を返す(ちょうど k 個のポップが完了)
2. n > k の場合:負の無限大(-math.inf)を返す(条件違反の経路)
3. i がスタックの総数と等しい場合:負の無限大を返す(使えるスタックが残っていない)
4. i が最後のスタック(len(stacks) - 1)の場合:
   • needed = k - n として、あと何個必要かを求める
   • needed がそのスタックの要素数より大きければ、負の無限大を返す
   • そうでなければ、そのスタックの末尾 needed 個の要素の合計を返す
5. それ以外の場合:
   • res = -math.inf、su = 0 で初期化する
   • sti をスタックの末尾から先頭へ向かってループさせながら:
     • su += stacks[i][sti](末尾からの累積和を更新)
     • localres = su + rec(i + 1, n + len(stacks[i]) - sti)
     • res = max(res, localres)
   • max(res, rec(i + 1, n)) を返す(このスタックを1つも使わないケースも含む)
6. メイン処理から rec(0, 0) を呼び出して結果を取得する

実装例(Python)

import math

class Solution:
    def solve(self, stacks, k):
        def rec(i, n):
            if n == k:
                return 0
            if n > k:
                return -math.inf
            if i == len(stacks):
                return -math.inf
            if i == len(stacks) - 1:
                needed = k - n
                if needed > len(stacks[i]):
                    return -math.inf
                else:
                    return sum(stacks[i][-needed:])
            res, su = -math.inf, 0
            for sti in range(len(stacks[i]) - 1, -1, -1):
                su += stacks[i][sti]
                localres = su + rec(i + 1, n + len(stacks[i]) - sti)
                res = max(res, localres)
            return max(res, rec(i + 1, n))

        return rec(0, 0)

ob = Solution()
stacks = [
    [50, -4, -15],
    [2],
    [6, 7, 8]
]
k = 4
print(ob.solve(stacks, k))

入力

[[50, -4, -15], [2], [6, 7, 8]], 4

出力

39

計算量と改善のポイント

この実装は、各スタックについて「何個ポップするか」を全パターン試すため、スタック数や要素数が増えると指数的に計算量が増大します。しかし、再帰の状態は (i, n) の組み合わせに限られるため、functools.lru_cache を使ってメモ化すれば、同一状態の再計算を避けられ、動的計画法として大幅に高速化できます。

また、負の値を含むスタックでは「全部取るのが最適とは限らない」点にも注意が必要です。上記のアルゴリズムは各スタックをスキップする選択肢(rec(i + 1, n))も評価しているため、このようなケースにも正しく対応できます。