Pythonで要素をポップして全スタックの合計を揃えるときの最大合計を求めるプログラム

複数のスタック（リスト）が与えられ、その中から任意のスタックを選んで任意の個数の要素をポップ（末尾から取り除く）できるものとします。このとき、すべてのスタックの合計値が等しくなるように操作を行った結果として達成できる最大の合計値を求めるのが本記事のテーマです。

たとえば、入力が stacks = [[3, 4, 5, 6], [5, 6, 1, 4, 4], [10, 2, 2, 2]] の場合、出力は 12 になります。具体的には次のような操作で実現できます。

• 1つ目のスタックから 6 をポップ → 残りは [3, 4, 5]、合計は 12
• 2つ目のスタックから 4, 4 をポップ → 残りは [5, 6, 1]、合計は 12
• 3つ目のスタックから 2, 2 をポップ → 残りは [10, 2]、合計は 12

解法の考え方：累積和とカウント

スタックから要素をポップすると、残るのは必ず「先頭からの連続した部分列」になります。つまり、各スタックで達成可能な合計値は、先頭から順に要素を足していった累積和のいずれかです。

そこで、すべてのスタックの累積和を辞書に記録し、それぞれの値が何個のスタックで現れたかをカウントします。スタックの総数と同じ回数だけ現れる累積和（＝全スタックで共通に作れる合計値）の中で最大のものが答えとなります。

アルゴリズムの手順

1. sums として空の辞書（defaultdict(int)）を用意する。
2. 各スタック stk について、変数 s を 0 で初期化し、各要素 n を順に加算しながら sums[s] をインクリメントする。
3. ans を 0 で初期化する。
4. sums の各ペア (s, f) について、f がスタック数以上 かつ s > ans であれば ans := s と更新する。
5. ans を返す。

要素の総数を N とすると、計算量は O(N) で済むため非常に効率的です。

Pythonでの実装例

from collections import defaultdict

class Solution:
    def solve(self, stacks):
        sums = defaultdict(int)
        # 各スタックの累積和をすべて記録する
        for stk in stacks:
            s = 0
            for n in stk:
                s += n
                sums[s] += 1

        # 全スタックで共通する累積和のうち最大のものを探す
        ans = 0
        for s, f in sums.items():
            if f >= len(stacks) and s > ans:
                ans = s
        return ans

ob1 = Solution()
stacks = [
    [3, 4, 5, 6],
    [5, 6, 1, 4, 4],
    [10, 2, 2, 2]
]
print(ob1.solve(stacks))

入力

stacks = [[3, 4, 5, 6], [5, 6, 1, 4, 4], [10, 2, 2, 2]]

出力

12

まとめ

この問題のポイントは、「ポップ後の残りの合計＝先頭からの累積和」であることに気づくことです。累積和の出現回数を辞書で管理すれば、全スタックで共通に作れる合計値を線形時間で見つけられます。また、どのスタックも空にできるため、共通の合計値が存在しない場合は初期値の 0 が返される仕組みになっています。