TopCoder SRM 546 Div1 Medium: FavouriteDigits

解法

上の桁から貪欲に埋めていく。

できるだけ小さい数字を入れてみて、残りの数字が条件を満たす最大のものだった時にN以上になるなら大丈夫。

コード

public class FavouriteDigits {

    public long findNext(long lN, int dSmall, int cSmall, int dBig, int cBig) {
	if (dSmall > dBig) {
	    return findNext(lN, dBig, cBig, dSmall, cSmall);
	}

	// 16桁に直す
	String N = String.format("%016d", lN);
	String ans = "";
	boolean leadingZero = true;// leading zeroの最中かどうか
	for (int pos = 0; pos < N.length(); pos++) {
	    for (int number = 0; number <= 9; number++) {
		int nextCSmall = cSmall;
		if ((dSmall != 0 || !leadingZero) && number == dSmall) {
		    // dsmallが0でleading zero中だったら
		    // 0を加えてもそれはカウントされない
		    nextCSmall--;
		}
		int nextBig = cBig;
		if (number == dBig) {
		    nextBig--;
		}
		if (isOK(ans + number, N, dSmall, nextCSmall, dBig, nextBig)) {
		    ans = ans + number;
		    cSmall = nextCSmall;
		    cBig = nextBig;
		    if (number != 0) {
			leadingZero = false;
		    }
		    break;
		}
	    }
	}
	return Long.parseLong(ans);
    }

    private boolean isOK(String prefix, String N, int dSmall, int cSmall,
	    int dBig, int cBig) {
	// prefixが決まっている状態で考えうる最大のsuffixを生成した場合、条件を満たせるかどうか調べる
	cSmall = cSmall < 0 ? 0 : cSmall;
	cBig = cBig < 0 ? 0 : cBig;

	int left = N.length() - prefix.length();
	if (cSmall + cBig > left) {
	    return false;
	}
	String suffix = "";

	for (int i = 0; i < cSmall; i++) {
	    suffix = dSmall + suffix;
	}
	for (int i = 0; i < cBig; i++) {
	    suffix = dBig + suffix;
	}
	while (suffix.length() < left) {
	    suffix = 9 + suffix;
	}

	return Long.parseLong(prefix + suffix) >= Long.parseLong(N);
    }
}