next_permutation の要領で自然数の分割を列挙することで \(O(N \cdot p(N))\) （ただし \(p(N)\) は自然数 \(N\) の分割数）時間で解くことができます。

自然数 \(N\) の分割 \(A_0 \ge A_1 \ge \dots \ge A_{l - 1} > 0\) に対して、\(A\) の辞書順後者（\(A\)より辞書順で大きい分割が存在すればその中で辞書順で最小のもの）を返し、存在しなければそれを報告するアルゴルズムを構築しましょう。

なお分割 \(A\) に対して適宜 \(A_{-1} = \infty\) と解釈することとします。

アルゴリズム

\(l \le 1\) のとき \(A\) の辞書順後者は存在しません。\(l \ge 2\) のときは、次の条件を満たす \(k\) が存在するので、そのようなもののうち最大のものを取ります。（これが \(A\) との間の LCP (longest common prefix) の長さになります。）

1. \(0 \le k \le l - 2\)（\(k = l - 1\) が許されていないのがポイント）
2. \(A_{k - 1} > A_k\)

このとき次のように定まる \(B\) が \(A\) の辞書順後者です。

\[ B_i = \begin{cases} A_i &(i < k) \\ A_i + 1 &(i = k) \\ 1 & ((k < i \le k + (A_k + \dots A_l) ) \end{cases} \]

証明

\(N\) の分割のうち \(A\) よりも辞書順で大きなもの全体の集合を \(\Sigma = A^\triangle\) とおきます。\(\mathrm{len}(A) \le 1\) のときは \(\Sigma = \emptyset\) で、\(\mathrm{len}(A) \ge 2\) ときは \(\Sigma \neq \emptyset\) です。以降後者の場合のみを考えます。

\(\Sigma\) においては \(A\) との LCP (longest common prefix) が長いほど辞書順で小さいことに着目し、次の 2 段階で考察しましょう。

• Step 1: \(A\) との LCP の長さの最大値 \(k\) を求める
• Step 2: \(A\) との LCP の長さが \(k\) である中で辞書順最小の分割を求める

Step 1

\(B = (B_0, \dots, B_{m - 1}) \in \Sigma, \ \mathrm{LCP}(A, B) = k\) と仮定すると、次が成り立ちます:

• \(B\) は自然数 \(N\) の分割である
• \((A_0, A_1, \dots, A_{k - 1}) = (B_0, B_1, \dots, B_{k - 1})\)
• \((A_k, A_{k + 1}, \dots, A_{l - 1}) < (B_k, B_{k + 1}, \dots, B_{m - 1})\)

最後の条件より \(A_k < B_k\) が成り立つので、\(A_{k + 1} + \dots + A_l > B_{k + 1} + \dots + B_{m - 1}\) となり、\(k \le l - 2\) が従います。また \(A_{k - 1} = B_{k - 1} \le B_k <A_k\) も従います。

以上より \(k\) は次の条件 (*) を満たす必要があります。

1. \(0 \le k \le l - 2\)
2. \(A_{k - 1} > A_k\)

さらに実は、逆にこれさえ満たせば必ず \(B \in \Sigma, \ \mathrm{LCP}(A, B) = k\) なる \(B\) が存在します。Step 2 でこれを構築することをもって、その証明としましょう。

Step 2

条件 (*) を満たす最大の \(k\) を取ると、\((A_k, A_{k + 1}, \dots, A_l)\) は次のいずれかの形の、長さ \(2\) 以上の分割です。

1. \((x, x, \dots, x)\)（ただし \(x > 0\))
2. \((x, x, \dots, x, r)\)（ただし \(x > r > 0\))

いずれも辞書順後者を持ち、それは \((x + 1, 1, 1, \dots, 1)\) の形で 表せます。

コード

Rust で書くと次のようになります。

fn next_partition(a: &mut Vec<usize>) -> bool {
    if a.len() <= 1 {
        return false;
    }
    let i = (1..a.len() - 1).rfind(|&i| a[i - 1] > a[i]).unwrap_or(0);
    let x = a[i] + 1;
    let y = a.drain(i..).sum::<usize>() - x;
    a.push(x);
    a.extend(std::iter::repeat(1).take(y));
    true
}

関連記事

next_permutation と同様に列挙できるもの一覧 - ブログ名

提出

Rust (122 ms): https://atcoder.jp/contests/abc226/submissions/64159026