山傘のプログラミング勉強日記

プログラミングに関する日記とどうでもよい雑記からなるブログです。

[AtCoder] AtCoder Beginner Contest 097 C - K-th Substring について

部分文字列に関する問題

C - K-th Substring

コンテスト中に解けず、ずっと考えたんですが答えが導けなかったので、解説を見て解くことができました。 解説を見た後では簡単な問題だと理解できたので、解法の発想に至るかどうかが大事な問題でした。

考え方

この問題の肝は、解答の文字列の候補の長さが  K で抑えられている点です。つまり、与えられた文字列の長さを  l とすると、部分文字列の先頭の文字の候補は  l となります。

文字列の長さが最大でも  K であることから、部分文字列の候補の数は、

 K \times (l - K + 1) + \dfrac{K(K - 1)}{2}

となります。文字列の  l - K +2 以降では、部分文字列の長さが  K より小さくなることに注意します。

 1 \leq K \leq 5

 1 \leq l \leq 5000

より、部分文字列の候補数は  25000 よりも小さくなることが分かります。

最後に、部分文字列を昇順にソートし、互いに異なる文字列をカウントして番号を振ることで、解答が得られます。

コード

import java.util.ArrayList;
import java.util.Collections;
import java.util.Scanner;
public class ProblemC {
    public static void main(String[] args) {
        Scanner scan = new Scanner(System.in);
        String s = scan.next();
        int l = s.length();
        int K = scan.nextInt();
        scan.close();
        ArrayList<String> list = new ArrayList<String>();
        for(int i = 0; i < l; i++) {
            for(int j = 1; j <= K; j++) {
                if(i + j <= l) {
                    list.add(s.substring(i, i + j));
                }else {
                    break;
                }
            }
        }
//     System.out.println(list.size());
//     int k = K * (l - K + 1) + K *(K - 1) / 2;
//     System.out.println(k);
        Collections.sort(list);
        String ans = list.get(0);
        if(K==1) {
            System.out.println(ans);
            System.exit(0);
        }
        int cnt = 1;
        for(int i = 0; i < list.size(); i++) {
            if(!list.get(i).equals(list.get(i + 1))) {
                cnt ++;
                if(cnt == K) {
                    ans = list.get(i + 1);
                    break;
                }
            }
        }
        System.out.println(ans);
    }
}