Typical DP Contest G: 辞書順 (動的計画法)

問題

G: 辞書順 - Typical DP Contest | AtCoder

解法

dp[i][c]:= i文字目以降の部分列で文字cから始まるものの個数、とする。i=N-1から順々にdpを埋めていき、埋め終わったらそこから文字列を復元すれば良い。

コード

import java.io.BufferedReader;
import java.io.InputStreamReader;

public class Main {

	public static void main(String[] args) throws Exception {
		BufferedReader r = new BufferedReader(new InputStreamReader(System.in), 1);

		char[] s = r.readLine().toCharArray();
		long K = Long.parseLong(r.readLine());

		int N = s.length;

		// dp[i][c]:= i文字目以降の部分列で文字cから始まるものの個数
		long[][] dp = new long[N + 1][26];
		for (int i = N - 1; i >= 0; i--) {
			for (int j = 0; j < 26; j++) {
				if (s[i] == (char) j + 'a') {
					dp[i][j] = 1;

					// i文字目以降の部分文字列の数は、i+1文字目以降の部分文字列の数の和+1である
					for (int k = 0; k < 26; k++) {
						dp[i][j] = Math.min(dp[i][j] + dp[i + 1][k], K + 1);
						if (dp[i][j] == K + 1) {
							break;
						}
					}
				} else {
					dp[i][j] = dp[i + 1][j];
				}
			}
		}

		long sum = 0;
		for (int i = 0; i < 26; i++) {
			sum = Math.min(sum + dp[0][i], K + 1);
		}
		if (K > sum) {
			System.out.println("Eel");
			return;
		}

		StringBuilder sb = new StringBuilder();
		for (int i = 0; i < dp.length; i++) {
			for (int c = 0; c < 26; c++) {
				if (K > dp[i][c]) {
					K -= dp[i][c];
				} else {
					while (s[i] != (char) c + 'a') {
						i++;
					}
					sb.append(s[i]);
					K--;// 部分文字列s[i]を出したので1つ減る
					break;
				}
			}
			if (K == 0) {
				break;
			}
		}
		System.out.println(sb.toString());

	}
}

大きいテストケースはBufferedReaderとInputStreamReaderを使わないと通らないなど……