计算数字总和到固定结果的组合数的算法难题

Question

这是我昨晚想到的一个谜题。我想出了一个解决方案，但效率不高，所以我想看看是否有更好的主意。

谜题是这样的:

给定正整数 N 和 T，你需要:

for i in [1, T], A[i] from { -1, 0, 1 }, such that SUM(A) == N

additionally, the prefix sum of A shall be [0, N], while when the prefix sum PSUM[A, t] == N, it's necessary to have for i in [t + 1, T], A[i] == 0

here prefix sum PSUM is defined to be: PSUM[A, t] = SUM(A[i] for i in [1, t])

the puzzle asks how many such A's exist given fixed N and T

例如，当N = 2，T = 4时，A的工作如下:

1 1 0 0
1 -1 1 1
0 1 1 0

但以下不要:

-1 1 1 1  # prefix sum -1
1 1 -1 1  # non-0 following a prefix sum == N
1 1 1 -1  # prefix sum > N

下面的 python 代码可以验证这样的规则，当给定 N 作为 expect 和 A 的实例作为 seq (有些人可能觉得阅读代码比阅读文字描述更容易) :

def verify(expect, seq):
    s = 0
    for j, i in enumerate(seq):
        s += i
        if s < 0:
            return False
        if s == expect:
            break
    else:
        return s == expect
    for k in range(j + 1, len(seq)):
        if seq[k] != 0:
            return False
    return True

我已经编写了我的解决方案，但它太慢了。以下是我的:

我将问题分解为两部分，一部分没有-1(只有{0, 1})，一部分有-1.

所以如果 SOLVE(N, T) 是正确答案，我定义一个函数 SOLVE'(N, T, B)，其中正 B 允许我将前缀和扩展到 [-B, N] 的区间而不是 [0, N]

所以实际上 SOLVE(N, T) == SOLVE'(N, T, 0)。

所以我很快意识到解决方案实际上是:

1. A 的前缀是一些有效的 {0, 1} 组合，具有正长度 l 和 o 1里面有
2. 在 l + 1 位置，我开始添加 1 个或多个 -1 并使用 B 跟踪数字。最大值将为 B + o 或取决于 A 中剩余的槽数，以较小者为准。
3. 递归调用SOLVE'(N, T, B)

在前面的 N = 2, T = 4 示例中，在其中一个搜索案例中，我会这样做:

1. 令 A 的前缀为 [1]，则我们有 A = [1, -, -, -]。
2. 开始添加-1。在这里我将只添加一个:A = [1, -1, -, -]。
3. 递归调用SOLVE'，这里我将调用SOLVE'(2, 2, 0) 来解决最后两个点。这里它只会返回 [1, 1]。然后其中一个组合产生 [1, -1, 1, 1]。

但是这个算法太慢了。

我想知道我该如何优化它或以任何不同的方式来看待这个可以提高性能的问题？(我只需要这个想法，而不是实现)

编辑:

一些示例将是:

T N RESOLVE(N, T)
3 2 3
4 2 7
5 2 15
6 2 31
7 2 63
8 2 127
9 2 255
10 2 511
11 2 1023
12 2 2047
13 2 4095
3 3 1
4 3 4
5 3 12
6 3 32
7 3 81
8 3 200
9 3 488
10 3 1184
11 3 2865
12 3 6924
13 3 16724
4 4 1
5 4 5
6 4 18

通常遵循指数时间解决方案(在 python 中):

import itertools
choices = [-1, 0, 1]
print len([l for l in itertools.product(*([choices] * t)) if verify(n, l)])

Answer 1

观察:假设 n 至少为 1，您所陈述问题的每个解决方案都以 [1, 0, .. ., 0]:即单个 1 后跟零个或多个 0。该点之前的解决方案部分是完全位于 [0, n-1] 中的步行，从 0 开始，在 n-1 结束，并且少于 t 个步骤。

因此，您可以将原来的问题简化为一个稍微简单的问题，即确定开始时 [0, n] 中有多少 t 步走在 0 并在 n 结束(其中每个步骤可以是 0、+1 或 -1 ，和以前一样)。

下面的代码解决了更简单的问题。它使用 lru_cache 装饰器来缓存中间结果；这是在 Python 3 的标准库中，或者有一个 recipe您可以下载适用于 Python 2 的版本。

from functools import lru_cache

@lru_cache()
def walks(k, n, t):
    """
    Return the number of length-t walks in [0, n]
    that start at 0 and end at k. Each step
    in the walk adds -1, 0 or 1 to the current total.

    Inputs should satisfy 0 <= k <= n and 0 <= t.
    """
    if t == 0:
        # If no steps allowed, we can only get to 0,
        # and then only in one way.
        return k == 0
    else:
        # Count the walks ending in 0.
        total = walks(k, n, t-1)
        if 0 < k:
            # ... plus the walks ending in 1.
            total += walks(k-1, n, t-1)
        if k < n:
            # ... plus the walks ending in -1.
            total += walks(k+1, n, t-1)
        return total

现在我们可以使用这个功能来解决您的问题。

def solve(n, t):
    """
    Find number of solutions to the original problem.
    """
    # All solutions stick at n once they get there.
    # Therefore it's enough to find all walks
    # that lie in [0, n-1] and take us to n-1 in
    # fewer than t steps.
    return sum(walks(n-1, n-1, i) for i in range(t))

solve(10, 100) 在我的机器上的结果和时间:

In [1]: solve(10, 100)
Out[1]: 250639233987229485923025924628548154758061157

In [2]: %timeit solve(10, 100)
1000 loops, best of 3: 964 µs per loop