Typical DP Contest G - 辞書順

http://tdpc.contest.atcoder.jp/tasks/tdpc_lexicographical

問題概要

英小文字からなる文字列Sの部分文字列のうち、辞書順でK番目のものを求めよ。ただしここでいう部分文字列は元の文字列で連続でないものも含む。また違うindexからなるものでも、結果として同じ文字列になるものはひとつとして扱う。

|S| <= 10⁶, K <= 10¹⁸

解法

dp(n)をn文字目を先頭としたときに作れるユニークな部分文字列の数とする。

next(α, i) = 文字列Sにおいて文字αがi番目以降初めて現れるインデックス

とすると、dp(n)は以下のように表せる。

これを求めれば「次に文字xを取ったら辞書順でいくつ進むか」がわかるので最終的な答えが求まる。

注意点（というか自分がハマったところ）として、|S|が大きいのでdpをメモ化再帰でやろうとするとスタックオーバーフローになること、またdpの値は64bit型整数でもオーバーフローしうることなどがあった。前者はdpをループで書き直し、後者はKを超える値は全部INFにまとめるという小細工で回避した。

感想

解説見ずにとおせてうれしかったです

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop;
 
immutable long INF = 10L ^^ 18 + 1000;
 
void main() {
    auto S = readln.chomp;
    auto N = S.length.to!int;
    auto K = readln.chomp.to!long;
 
    auto chars = new int[][](26, N);
    foreach (i; 0..26) {
        int tmp = -1;
        for (int j = N - 1; j >= 0; j--) {
            if (S[j] - 'a' == i) tmp = j;
            chars[i][j] = tmp;
        }
    }
 
    auto dp = new long[](N);
    fill(dp, 1);
 
    for (int i = N - 2; i >= 0; i--) {
        foreach (j; 0..26) {
            int m = chars[j][i + 1];
            if (m == -1) continue;
            if (dp[m] == INF) {
                dp[i] = INF;
                continue;
            }
            dp[i] += dp[m];
            if (dp[i] > INF || dp[i] <= 0) {
                dp[i] = INF;
                continue;
            }
        }
    }
 
 
    string ans = "";
 
 
    void dfs(int p, long acm) {
        if (acm == K) return;
        foreach (i; 0..26) {
            int m = chars[i][p];
            if (m == -1) continue;
            if (acm + dp[chars[i][m]] >= K) {
                ans ~= i.to!char + 'a';
                dfs(chars[i][m] + 1, acm + 1);
                return;
            }
            acm += dp[chars[i][m]];
        }
    }
 
 
    dfs(0, 0);
    if (ans == "") writeln("Eel");
    else writeln(ans);
}