Codeforces Round #Pi Div2 D: One-Dimensional Battle Ships

問題

codeforces.com

解法

残ったボードに存在しうる船の数を常に保持するようにする。

コード

import java.util.Scanner;
import java.util.TreeMap;

public class Main {

	public void solve() {
		Scanner scanner = new Scanner(System.in);
		int boardLength = scanner.nextInt();
		int shipNum = scanner.nextInt();
		int shipLength = scanner.nextInt();

		// (ボードの最初のマス, ボードの長さ)
		TreeMap<Integer, Integer> treeMap = new TreeMap<>();
		treeMap.put(1, boardLength);

		// 最大でボード上に載り得る船の数
		int maxHold = (boardLength + 1) / (shipLength + 1);

		int Q = scanner.nextInt();
		for (int q = 0; q < Q; q++) {
			int shot = scanner.nextInt();

			// 撃ったマスの直前に残っているボードを取り出す
			Integer hitKey = treeMap.floorKey(shot);
			if (hitKey == null) {
				// 撃ったマスの直前に残っていなければスルー
				continue;
			}

			int hitLength = treeMap.get(hitKey);
			if (hitKey + hitLength - 1 < shot) {
				// 取り出したボードに当たっていなければスルー
				continue;
			}

			// 撃たれたボードに載る船の最大数
			int beforeMax = (hitLength + 1) / (shipLength + 1);

			// 分割されて新しく出来たボード
			int nextLenght1 = shot - hitKey;
			int nextLenght2 = hitLength - 1 - nextLenght1;
			if (hitKey == shot || nextLenght1 < shipLength) {
				treeMap.remove(hitKey);
			}
			if (nextLenght1 >= shipLength) {
				treeMap.put(hitKey, nextLenght1);
			}
			if (nextLenght2 >= shipLength) {
				treeMap.put(shot + 1, nextLenght2);
			}

			// 撃たれた後に載る船の最大数を計算する
			int afterMax = (nextLenght1 + 1) / (shipLength + 1)
					+ (nextLenght2 + 1) / (shipLength + 1);
			maxHold -= (beforeMax - afterMax);

			// 載りうる数が船の数より少なくなれば必ず当たるようになる
			if (maxHold < shipNum) {
				System.out.println(q + 1);
				scanner.close();
				return;
			}
		}
		System.out.println(-1);
		scanner.close();

	}

	public static void main(String[] args) {
		new Main().solve();
	}

}