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;
}