Atcoder - 161D 数位dp

<Atcoder - 161D> 数位dp

https://atcoder.jp/contests/abc161/tasks/abc161_d

题意：

定义一种数 lun number：该数任意相邻的两个数位差的绝对值不超过1

比如 123，234，12121，1010 都是 lun number；而 256，18，102 都不是 lun number

输入1e5以内的 k ，输出第 k 小的 lun number 的值

思路：

设dp[i][j]表示最高位是i, 位数是j的lun number的个数
根据极限样例, j最多为10位, i在0 ~ 9
首先初始化所有dp[i][1] = 1 (以i开头的一位数就是i本身, 只有1个)
转移方程3种情况:
1) dp[i][j] = dp[i][j - 1] + dp[i + 1][j - 1];
//最高位i = 0时的j位数只能以i和i + 1作最高位的j - 1位数作为后缀
2) dp[i][j] = dp[i - 1][j - 1] + dp[i][j - 1];
//最高位i = 9时的j位数只能以i - 1和i作最高位的j - 1位数作为后缀
3) dp[i][j] = dp[i - 1][j - 1] + dp[i][j - 1] + dp[i + 1][j - 1]; 
//最高位非0非9时, 以i - 1、i本身、i + 1作最高位的j - 1位数均可作为最高位是i的j位数的后缀

接下来进行构造:
首先定位第k小的lun number的位数bit
统计所有i位数中有多少lun number, 做一遍前缀和存进ans[i]
从低位向高位遍历, 找到第一个大于等于k的位数即为k的位数bit

确定好位数之后, 开始枚举每一位是什么
如果是第一位, 我们暴力枚举1 ~ 9
对k执行递减, 看k处于谁开头的bit位数
用找定位位数的方法, 我们可以定位第一个位置的值(已经记录了以i开头, 有j位的lun number有多少个)
然后后面的位置的数只有最多三种选择: 上次减一; 上次同样; 上次加一(在0 ~ 9之间)
按顺序去搜索alter - 1 < alter < alter + 1
不断定位当前位的值, 定位好了就dfs进入下一层
一直记忆化搜所即可

AC代码：

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
LL dp[10][12]; //dp[i][j]代表以i开头的j位数有多少个lun number 
LL k, ans[12]; //ans[i]表示i位lun number的个数
int bit;

void dfs(int num, int alter) { //第一个参数表示第几位(层数), 第二个参数表示变化方式(-1/不变/+1)
	if(num > bit) return ;
	if(num == 1) { //第一位直接枚举1 ~ 9
		for(int i = 1; i <= 9; i++) { //后缀共bit - num + 1位
			if(k > dp[i][bit - num + 1]) k -= dp[i][bit - num + 1];
			else {
				cout << i; //最高位的值
				dfs(num + 1, i); //这时候dfs进入下一层看第二高位是几, 进入else特判内, 记忆化搜索
				break;
			}
		}
	}
	else {
		for(int i = alter - 1; i <= alter + 1; i++) {
			if(!(i >= 0 && i <= 9)) continue;
			if(k > dp[i][bit - num + 1]) k -= dp[i][bit - num + 1];
			else {
				cout << i;
				dfs(num + 1, i);
				break;
			}
		}
	}
}

int main() {
	cin >> k;
	for(int i = 0; i < 10; i++) dp[i][1] = 1;
	ans[1] = 9LL;
	for(int j = 2; j <= 10; j++) { //位数:2 ~ 10
		for(int i = 0; i < 10; i++) { //最高位:0 ~ 9
			if(i == 0) dp[i][j] = dp[i][j - 1] + dp[i + 1][j - 1];
			//最高位i = 0时的j位数只能以i和i + 1作最高位的j - 1位数作为后缀
			else if(i == 9) dp[i][j] = dp[i - 1][j - 1] + dp[i][j - 1];
			//最高位i = 9时的j位数只能以i - 1和i作最高位的j - 1位数作为后缀
			else dp[i][j] = dp[i - 1][j - 1] + dp[i][j - 1] + dp[i + 1][j - 1]; 
			//最高位非0非9时, 以i - 1、i本身、i + 1作最高位的j - 1位数均可作为最高位是i的j位数的后缀
			if(i) ans[j] += dp[i][j]; //没有前导0时就将dp[i][j]表示的个数存入ans[j]
		}
	}
	for(int i = 1; i <= 20; i++) {
		if(k > ans[i]) k -= ans[i];
		else {
			bit = i; //所求转化为:位数为bit位的所有数中的第k小的那个(将当前k定位在bit位)
			break;
		}
	}
	dfs(1, 0);
	cout << endl;
	return 0;
}