EDIT: 最初に投稿した時は、N個の整数の合計の一般解を求める問題だと思っていたのですが、一般解だとNP完全問題になってしまうので、出題者の意図は３つ限定だったのだろうと思います。ですので３つの合計に改題しました。回答を考えてくださった皆さんすみません！

ある任意の長さのリストに符号つきの整数が格納されています。値の重複を許します。そこから３つの整数を取り出すとき、合計が0になる組み合わせを求めなさいという問題が解けません。どうぞご教授ください。

条件

  • itertoolsを使えば下のように簡単に計算できますが、itertoolsは使わないものとします。
  • メモリを3回走査するとtime complexity(CPUタイム)はO(N^3)になります。それよりは効率のよい方法で実装したいです。
  • space complexity(メモリ効率)は犠牲にしていいです。
  • ソートやリストのコピーなど前処理を行っても構いません。

検算用のコード

import itertools
l=[-3, 2, 1, -2, 0, 5, -1 , -1]
for c  in itertools.combinations(l, 3):
    if sum(c) == 0:
        print c

答え

(-3, 2, 1)
(-3, -2, 5)
(2, -2, 0)
(2, -1, -1)
(1, 0, -1)
(1, 0, -1)

今まで調べたこと

itertools.combinationのマニュアルには同等のアルゴリズムのソースがありました。このアルゴリズムはO(N)だそうです。

def combinations(iterable, r):
    # combinations('ABCD', 2) --> AB AC AD BC BD CD
    # combinations(range(4), 3) --> 012 013 023 123
    pool = tuple(iterable)
    n = len(pool)
    if r > n:
        return
    indices = range(r)
    yield tuple(pool[i] for i in indices)
    while True:
        for i in reversed(range(r)):
            if indices[i] != i + n - r:
                break
        else:
            return
        indices[i] += 1
        for j in range(i+1, r):
            indices[j] = indices[j-1] + 1
        yield tuple(pool[i] for i in indices)

これだけでも十分速いと思いますが、与えられた問題とデータの性質に合わせてより効率的な方法はないかと模索しています。

条件のところに書いていますが、ヒントはソートしても良いし、リストのコピーを作っても良い、とのことです。また取り出す個数が3でなくて2ならどうか？とも聞かれました。2つなら辞書にd[-1]=1などとして、リストに存在する符号が反対の値の組合せを保存しておけば、O(1)で計算できることになります。しかしこれでは重複する値に対応できないし、3つ以上に拡張できるのか、全く見当もつきません。

類題

部分和問題という類題は任意の数の部分和を探すもので、僕に与えられた問題はサブセットということになりそうです。なんとなく、見た目より遥かに難しい問題だということはわかりました。

部分和問題。NP完全だそうです...。

さらに調べてわかったこと

取り出す個数が３の時は"3sum問題"として、ダイナミックプログラミング(DP、動的計画法)のO(N^2)の解が知られている。

本家StackOverflow N=3の時の回答　コードも簡単。
３和問題