题目链接：https://atcoder.jp/contests/abc200/tasks/abc200_e
E - Patisserie ABC 2
题意
(n^3) 个三元组 ((x,y,z) (1 le x,y,z le n)) 按照以下三个关键字从小到大排序：
- (x + y + z)
- (x)
- (y)
计算第 (k) 个三元组的值。
题解
由于数位和的值为第一关键字，所以可以先推出第 (k) 个三元组的数位和。
数位和为 (s) 的三元组个数相当于把 (s) 个 (1) 分为 (3) 堆，且每堆 (1) 的总数不超过 (n) 。
不妨更一般性地，考虑：
将 (s) 个无差别小球分为 (m) 个非空堆，且每堆小球的总数不超过 (n) 的总情况数：
假如已有 (i) 堆个数超过 (n) ，那么剩余 (s - i imes n) 个小球，将这些小球分为 (m) 堆的情况数为 (C_{m}^{i} imes C_{s - i imes n - 1}^{m - 1}) ，意义为先从 (m) 堆中选出 (i) 堆各放置 (n) 个球，然后再在剩下 (s - i imes n - 1) 个空隙中插入 (m - 1) 个隔板。
但在剩余 (s - i imes n) 小球的 (C_{s - i imes n - 1}^{m - 1}) 种分法中仍可能存在分得的堆中小球个数大于 (n) 的情况，即 (i) 堆的情况数是包含 (i+1, i+2,dots, m) 堆的，那么由容斥原理可以推得，总情况数为： (sum limits _{i = 0}^{m} (-1)^i imes C_{m}^{i} imes C_{s - i imes n - 1}^{m - 1}) 。
由此推得数位和 (digit\_sum) ，之后枚举每一位的值，比如第一位为 (i) 的情况数可以看作将 (digit\_sum - i) 个小球再分为 (2) 堆，且每堆小球的总数不超过 (n) 。若 (k) 大于当前位当前值的情况数则减去，否则递归枚举下一位即可。
代码
#include <bits/stdc++.h>
using namespace std;
#define int long long
constexpr int digit_num = 3; //数位个数
inline int C(int n, int m) { //从 n 个小球中选出 m 个，无序
if (n < 0 or m > n) {
return 0;
}
if (m == 0) {
return 1;
} else if (m == 1) {
return n;
} else if (m == 2) {
return n * (n - 1) / 2;
} else if (m == 3) {
return n * (n - 1) * (n - 2) / 6;
} else { //本题 m 的值不会超过数位个数 3
return 0;
}
}
inline int n_to_m_pieces(int n, int m, int lim) { //将 n 个小球分为 m 堆，且每堆小球个数不超过 lim
int res = 0;
for (int i = 0; i <= m; i++) {
res += (i & 1 ? -1 : 1) * C(m, i) * C(n - i * lim - 1, m - 1); //容斥原理
}
return res;
}
inline void dfs(int n, int k, int left, int digit_sum) { //从前往后推导每一位
if (left == 0) { //之后无剩余位
cout << digit_sum << "
";
return;
}
for (int i = 1; i <= n; i++) {
if (digit_sum - i > left * n) { //当前数位值过小
continue;
}
if (k > n_to_m_pieces(digit_sum - i, left, n)) { //当前数位值为 i 的情况数
k -= n_to_m_pieces(digit_sum - i, left, n);
} else {
cout << i << ' ';
dfs(n, k, left - 1, digit_sum - i); //推导下一位
break;
}
}
}
inline int get_digit_sum(int n, int& k) { //推出三个数位的和
for (int i = digit_num; ; i++) { //数位和至少为数位个数
if (k > n_to_m_pieces(i, digit_num, n)) {
k -= n_to_m_pieces(i, digit_num, n); //数位和为 i 的情况数
} else {
return i;
}
}
}
signed main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n, k;
cin >> n >> k;
dfs(n, k, digit_num - 1, get_digit_sum(n, k));
return 0;
}