【Java】AtCoder Beginner Contest 217のD問題をTreeMapを使って解く

概要

AtCoder Beginner Contest 217のD問題 について、私はJavaで解いた際に、TreeMapの機能を使いました。あまり普段使わない機能だったので、今回記事として残しておきます。
なお、D問題の概要は、一つの線をどんどん分割していって、指定された地点はどこの線に含まれるか、というのを出力するというものです。

考察

指定された地点がどこの線に含まれるかを高速に取得することができれば、この問題はほぼ解けたも同然です。解説ではC++でlower_boundという関数で周辺値を取得していますが、これと同じようなことをJavaで実現できれば良いです。

対応

まず前提として、分割された線をTreeMapで管理します。keyが線の始まりの地点で、valueに線の長さをセットします。
そして、TreeMapのfloorEntryを使用すれば解説と似たようなことが実現できます。floorEntryはこちらのstackoverflowの記事にある通り、指定したkey以下の値で最大のEntryを取得します。今回の問題だと、指定した地点が含まれる線のEntryが取得できます。

実装サンプル

こちらの私の提出した回答があります。
線を分割する部分のロジックだけ、以下に紹介します。

public class Main {
    static TreeMap<Integer, Integer> rangeTreemap = new TreeMap<>();
   　  ・
   　  ・
   　  ・
    // 分割する地点を引数
    private static void split(int x) {
        // indexは0始まりで管理しているのでマイナスする
        int targetIndex = x - 1;
	// TreeMapが空の場合は初期状態とみなす
        if (rangeTreemap.isEmpty()) {
            rangeTreemap.put(0, x);
            rangeTreemap.put(x, l - x);
        } else {
	    // ここでその地点が含まれるEntryを取得
            Map.Entry<Integer, Integer> entry = rangeTreemap.floorEntry(targetIndex);
            // 分割した結果を算出しTreeMapを更新
            int before = targetIndex - entry.getKey() + 1;
            int after = entry.getValue() - before;
            rangeTreemap.put(entry.getKey(), before);
            rangeTreemap.put(x, after);
        }
    }
}