数値のリストの中身を足し合わせた結果が特定の値になったときの処理

Question

python3を使用しています。

例えば数値のリストlst[2, 1, 1, 2, 1, 3, 2, 1, 3, 2, 1, 4, 1, 4, 1, 2, 1, 3, 1, 4, 1, 3, 1, 2, 1, 6, 2, 1, 2, 1, 2, 1, 2, 1,]
があったとします。
このとき左から足していって3,6,4の並びになった時にTrueを返すようなプログラムを実現したいです。すこし特殊な状況なので人力で処理を書きます。
この場合だと、

2+1=3...第1段階クリア(lst[0]+lst[1]==3)
1+2+1+3>6...6を超えているのでNG(lst[2]+lst[3]+lst[4]+lst[5])
次はlst[1]からやり直し
1+1+2>3...3を超えているのでNG(lst[1]+lst[2]+lst[3])
次はlst[2]から
1+2=3...第1段階クリア(lst[2]+lst[3]==3)
1+3+2=6...第2段階クリア(lst[4]+lst[5]+lst[6]==6)
1+3=4...第3段階クリア(lst[7]+lst[8]==4)
この時Trueを返す(インデックスも)
次はlst[3]から以下同様最後まで

以上のようなプログラムです。
単に3+6+4で左からまとめて足して13になっていればよいというわけではなく、必ず3,6,4の区切りである必要があります。
自分なりに少し考えてみたのですが、ループ処理中の配列のインデックス指定や、失敗した際に次の要素からやり直すこと、ループ回数が増えて処理が重くなる、同じ計算をしている？等々色々な要素が入り混じっていてよくわからなくなってしまいました。

python、アルゴリズム共に初学者で余り明るくないのでご教授いただければと思います。

Answer 1

数値のリストに正の数しか含まれない前提で良ければ、累積和を使ったアプローチが使えると思います。

まず、数値リスト(lst)の先頭から累積和を計算して、辞書の形(acc)で用意しておきます。
lstの先頭から順に足していって 3, 6, 4 の並びになるということは、累積和で 3, 9, 13 パタンが出現するということになります。
acc中に 3, 9, 13 が全てあれば、lst[0]から足していったときに 3, 6, 4のパタンがあったことになります。

lst[1]から足していったときに 3, 6, 4のパタンがあるかをチェックするために
累積和の辞書(acc)を作り直すと時間がかかってしまうため、
代わりに 3+lst[0], 9+lst[0], 13+lst[0] が出現するかでチェックします。

同様に、lst[n]から足していったときに 3, 6, 4 の並びがあるかは、
3, 9, 13 それぞれにsum(lst[0:n])を足した数字があるかでチェックできます。
この3, 9, 13に足す数値をoffsetで管理することで、累積和やsumを何度も計算するのを避けることができます。

from itertools import accumulate

# 以下のように0や負の数が含まれる場合には、このプログラムは期待通りに動作しません
# lst = [1, 2, 1, 3, 2, 1, 3, 0]
# lst = [13, -4, -5, -1]
def solve(lst, targets):
    acc = {a: i for i, a in enumerate(accumulate(lst))}
    numbers = list(accumulate(targets))
    offset = 0
    res = []
    for i, l in enumerate(lst):
        if all(n+offset in acc for n in numbers):
            res.append((i, acc[numbers[-1]+offset]+1)) # 3, 6, 4 になる部分の始点と終点を記録
        offset += l
    return res


lst = [2, 1, 1, 2, 1, 3, 2, 1, 3, 2, 1, 4, 1, 4, 1, 2, 1, 3, 1, 4, 1, 3, 1, 2, 1, 6, 2, 1, 2, 1, 2, 1, 2, 1]
res = solve(lst, (3, 6, 4))
for l, r in res:
    print(f'lst[{l}:{r}] => {lst[l:r]}')
# lst[2:9] => [1, 2, 1, 3, 2, 1, 3]
# lst[6:12] => [2, 1, 3, 2, 1, 4]
# lst[17:23] => [3, 1, 4, 1, 3, 1]


# (3, 6, 4)以外のパタンも探せます
res = solve(lst, (3, 8))
for l, r in res:
    print(f'lst[{l}:{r}] => {lst[l:r]}')
# lst[5:10] => [3, 2, 1, 3, 2]
# lst[8:13] => [3, 2, 1, 4, 1]
# lst[15:20] => [2, 1, 3, 1, 4]
# lst[23:27] => [2, 1, 6, 2]
# lst[26:33] => [2, 1, 2, 1, 2, 1, 2]

Answer 2

プログラムは少し長くなりますが、パターンを使うことで高速にできる可能性があります。
このアルゴリズムをnumpyのndarrayでやるとかなり速いと思います。

#パターン作成

def enhance(l):
    lcopy = list(l)
    lcopy.append(1)
    return lcopy

def increase(l):
    lcopy = list(l)
    lcopy[-1] += 1
    return lcopy

def make_pattern(n):
    if n == 1:
        return [[1]]
    else:
        spattern = make_pattern(n-1)
        return [increase(l) for l in spattern] + [enhance(l) for l in spattern]

pattern3 = make_pattern(3)
pattern4 = make_pattern(4)
pattern6 = make_pattern(6)

#パターン作成終了

def check(lst, pattern):
    for p in pattern:
        if lst[:len(p)] == p:
            return len(p)
    return False

def check_head(lst):
    result3 = check(lst, pattern3)
    if result3 == False:
        return False
    else:
        pass
    result6 = check(lst[result3:], pattern6)
    if result6 == False:
        return False
    else:
        pass
    result4 = check(lst[result3+result6:], pattern4)
    if result6 == False:
        return False
    else:
        return (result3, result3+result6, result3+result6+result4)

def check_all(lst):
    for i in range(len(lst)):
        result_head = check_head(lst[i:])
        if result_head != False:
            return (result_head[0]+i, result_head[1]+i, result_head[2]+i)
    return False

実行すると、

>>> lst = [2, 1, 1, 2, 1, 3, 2, 1, 3, 2, 1, 4, 1, 4, 1, 2, 1, 3, 1, 4, 1, 3, 1, 2, 1, 6, 2, 1, 2, 1, 2, 1, 2, 1,]
>>>
>>> check_all(lst)
(4, 7, 9)

ちなみに、パターンを表示すると以下のようになります。

>>> print(pattern3)
[[3], [1, 2], [2, 1], [1, 1, 1]]
>>> print(pattern4)
[[4], [1, 3], [2, 2], [1, 1, 2], [3, 1], [1, 2, 1], [2, 1, 1], [1, 1, 1, 1]]
>>> print(pattern6)
[[6], [1, 5], [2, 4], [1, 1, 4], [3, 3], [1, 2, 3], [2, 1, 3], [1, 1, 1, 3], [4, 2], [1, 3, 2], [2, 2, 2], [1, 1, 2, 2], [3, 1, 2], [1, 2, 1, 2], [2, 1, 1, 2], [1, 1, 1, 1, 2], [5, 1], [1, 4, 1], [2, 3, 1], [1, 1, 3, 1], [3, 2, 1], [1, 2, 2, 1], [2, 1, 2, 1], [1, 1, 1, 2, 1], [4, 1, 1], [1, 3, 1, 1], [2, 2, 1, 1], [1, 1, 2, 1, 1], [3, 1, 1, 1], [1, 2, 1, 1, 1], [2, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]]

Answer 3

条件を満たす部分リストを列挙するプログラムをかいてみました。
列挙結果の個数が 0 かどうかで True/False 判定するればよいだけなので。

p.py

# 連続する要素の部分和が n になる場所を列挙する
def sublist(lst, n):
  ans = []
  for i in range(len(lst)):
    for length in range(1, len(lst) - i):
      j = i + length - 1
      if (sum(lst[i:j])) == n:
        ans.append([lst[i:j], lst[j:]])
  return ans

# 先頭の連続する要素の和が n になるものを列挙する
def headlist(lst, n):
  ans = []
  for i in range(len(lst)):
    if sum(lst[0:i]) == n:
      ans.append([lst[0:i], lst[i:]])
  return ans

lst = [
  2, 1, 1, 2, 1, 3, 2, 1, 3, 2,
  1, 4, 1, 4, 1, 2, 1, 3, 1, 4,
  1, 3, 1, 0, 2, -2,1, 6, 2, 1, 2, 1,
  2, 1, 2, 1
]

ans1 = sublist(lst, 3)
ans2 = []
for x in ans1:
  for y in headlist(x[1], 6):
    ans2.append([x[0] + y[0], y[1]])

ans3 = []
for x in ans2:
  for y in headlist(x[1], 4):
    ans3.append([x[0] + y[0], y[1]])

for z in ans3:
  print(z[0])

質問文にあった数列に 0, 2, -2 を挿入しています。
実行結果では、これを含む 部分リストが列挙されることが確認できています。

実行例：

ここでは同じ計算を避ける代わりに、大量のメモリをつかっています。
２番目以降の和のチェックをする部分は再帰をつかって書くこともできるかもしれません。
（そうすれば、和のパターンを配列で渡してチェックプログラムを書くことができそうです）

Answer 4

処理を分ければ見通しが良くなります。

def checksub2(lst,x):
# リストの先頭から幾つか足すとその数になるか？
  s = 0
  for i, v in enumerate(lst):
    s += v
    if s == x:
      return i+1
    if s > x:
      return False
  return False

def checksub(lst):
# リストの先頭からが条件を満たすか？
  return (
    (n1 := checksub2(lst,3)) and
    (n2 := checksub2(lst[n1:],6)) and
           checksub2(lst[n1+n2:],4)
  )

def check(lst):
  for i in range(0,len(lst)):
    if checksub(lst[i:]):
      return i
  return False

print(check([2, 1, 1, 2, 1, 3, 2, 1, 3, 2, 1, 4, 1, 4, 1, 2, 1, 3, 1, 4, 1, 3, 1, 2, 1, 6, 2, 1, 2, 1, 2, 1, 2, 1,]))

Python3.8以降でない場合。

def checksub(lst):
  n1 = checksub2(lst,3)
  if n1:
    n2 = checksub2(lst[n1:],6)
    if n2:
      return checksub2(lst[n1+n2:],4)
  return False

⇒ 代入式無いとたるい感じですね。

#追記
求めるパターン(3,6,4)がコード中に埋め込まれているのは、やな感じなので、引き数で渡すことにします。

def checksub2(lst,x):
  s = 0
  for i, v in enumerate(lst):
    s += v
    if s == x:
      return i+1
    if s > x:
      return False
  return False

def checksub(lst, pattern):
  i = 0
  for n in pattern:
    if not (j := checksub2(lst[i:],n)):
      return False
    i += j
  return True

def check(lst, pattern):
  for i in range(0,len(lst)):
    if checksub(lst[i:],pattern):
      return i
  return False

print(check([2, 1, 1, 2, 1, 3, 2, 1, 3, 2, 1, 4, 1, 4, 1, 2, 1, 3, 1, 4, 1, 3, 1, 2, 1, 6, 2, 1, 2, 1, 2, 1, 2, 1,],
      [3,6,4]))

分かりやすさ優先で書いてますので、速度を求めるならスライスを渡さず、「リスト全体と開始添え字」を渡すといいです。