公式

D - Coefficient Stair 解説 by en_translator 印刷ページ


To exhaustively enumerate objects lexicographically, DFS (Depth-First Search) is effective.

One can consider the following recursive function to lexicographically enumerate all combinations for the first \(i\) terms \(A _ 1,A _ 2,\ldots,A _ i\) satisfying the condition:

  1. If \(i=N\) (that is, if we have already determined the first \(N\) terms):
    1. If \(\displaystyle\sum _ {j=1} ^ Nj\times A _ j=K\), print \(A\).
    2. Otherwise, print nothing.
  2. Otherwise, try all possible integers as the \((i+1)\)-th term, from \(0\) through \(\Biggl\lfloor\dfrac{K-\sum _ {j=1} ^ ij\times A _ j}{i+1}\Biggr\rfloor\), and advance recursively.

In Python, the function can be expressed as follows:

def dfs(i, A):
    s = sum((j + 1) * a for j, a in enumerate(A[:i]))
    if i == N:
        if s == K:
            print(*A)
        return
    for x in range(0, (K - s) // (i + 1) + 1):
        A[i] = x
        dfs(i + 1, A)

Calling this function with dfs(0, [0] * N) would produce the correct results, but meeting the execution time limit is difficult on its own.

The function is called with \(i=N\) for each length-\(N\) sequence \(A\) of non-negative integers such that \(\displaystyle\sum _ {j=1} ^ Nj\times A _ j\le K\). Let \(f(N,K)\) be the number of length-\(N\) sequences \(A\) of non-negative integers such that \(\displaystyle\sum _ {j=1} ^ Nj\times A _ j=K\). This equals \(\displaystyle\sum _ {k=0} ^ Kf(N,k)\), which can grow up to \(10000100000\) under the constraints of this problem.

Here, modify the recursive function as follows:

  1. If \(i=N-1\) (that is, if we have already determined the first \((N-1)\) terms):
    1. If we can appropriately set \(A _ N\) so that \(\displaystyle\sum _ {j=1} ^ Nj\times A _ j=K\), print \(A\).
    2. Otherwise, print nothing.
  2. Otherwise, try all possible integers as the \((i+1)\)-th term, from \(0\) through \(\Biggl\lfloor\dfrac{K-\sum _ {j=1} ^ ij\times A _ j}{i+1}\Biggr\rfloor\), and advance recursively.

In Python, the function can be expressed as follows:

def dfs(i, A):
    s = sum((j + 1) * a for j, a in enumerate(A[:i]))
    if i == N - 1:
        if (K - s) % N == 0:
            A[i] = (K - s) // N
            print(*A)
        return
    for x in range(0, (K - s) // (i + 1) + 1):
        A[i] = x
        dfs(i + 1, A)

Now we can prove that the function with \(i=N-1\) is called \(O(Nf(N,K))\) times. Hence, the function is called a total of \(O(N ^ 2f(N,K))\) times, while we spend \(O(N)\) time to compute s, so this algorithm runs in a total of \(O(N ^ 3f(N,K))\) time. In some languages, this algorithm will meet the execution time limit on its own. By passing s as an argument of the recursive function, the time complexity can be reduced to \(O(N ^ 2f(N,K))\), which is now fast enough.

The following is sample code.

#include <iostream>
#include <vector>

int main() {
    using namespace std;
    int N, K;
    cin >> N >> K;

    vector<int> path(N);
    [N, K](this auto self, int index, int sum, vector<int>& path) -> void {
        if (index + 1 == N) { // When we determine the last one term,
            if ((K - sum) % N == 0) { // only if we can make the sum K,
                path.back() = (K - sum) / N;
                for (int p : path) // print it.
                    cout << p << ' ';
                cout << '\n';
            }
            return;
        } 

        // Decide the i-th term, and call the function recursively
        for (int i = 0; sum + i * (index + 1) <= K; ++i) {
            path[index] = i;
            self(index + 1, sum + i * (index + 1), path);
        }
    }(0, 0, path);
    
    return 0;
}
N, K = map(int, input().split())

def dfs(i, s, A):
    if i == N - 1: # When we determine the last one term,
        if (K - s) % N == 0: # only if we can make the sum K,
            A[i] = (K - s) // N
            print(' '.join(str(a) for a in A)) # print it.
        return

    # Decide the i-th term, and call the function recursively
    for x in range(0, (K - s) // (i + 1) + 1):
        A[i] = x
        dfs(i + 1, s + x * (i + 1), A)

dfs(0, 0, [0] * N)

投稿日時:
最終更新: