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:
- If \(i=N\) (that is, if we have already determined the first \(N\) terms):
- If \(\displaystyle\sum _ {j=1} ^ Nj\times A _ j=K\), print \(A\).
- Otherwise, print nothing.
- 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:
- If \(i=N-1\) (that is, if we have already determined the first \((N-1)\) terms):
- If we can appropriately set \(A _ N\) so that \(\displaystyle\sum _ {j=1} ^ Nj\times A _ j=K\), print \(A\).
- Otherwise, print nothing.
- 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)
投稿日時:
最終更新: