SRM554 Div2 Easy Practice

TreeSetの使い方を学ぶがてら、解いてみた。最初のコードは自分で書いたやつ。二番目は正解者のコードを参考にreviseしたやつ

import java.util.*;

public class TheBrickTowerEasyDivTwo {

	public int find(int redCount, int redHeight, int blueCount, int blueHeight) {
		int redtemp = redCount;
		int bluetemp = blueCount;
		TreeSet<Integer> set = new TreeSet<Integer>();
		set.add(redHeight);
		set.add(blueHeight);

		int temp = redHeight; // red ones are at the bottom
		redCount--;
		while (redCount != 0 || blueCount != 0) {
			if (blueCount > 0) {
				temp += blueHeight;
				blueCount--;
				set.add(temp);
			} else
				break;
			if (redCount > 0) {
				temp += redHeight;
				redCount--;
				set.add(temp);
			} else
				break;
		}
		
		redCount = redtemp;
		blueCount = bluetemp;
		
		temp = blueHeight; // blue ones are at the bottom
		blueCount--;
		while (redCount != 0 || blueCount != 0) {
			if (redCount > 0) {
				temp += redHeight;
				redCount--;
				set.add(temp);
			} else
				break;
			if (blueCount > 0) {
				temp += blueHeight;
				blueCount--;
				set.add(temp);
			} else 
				break;
		}
		return set.size();
	}

}

長いですね。reviseしましょう。
まず考え方ですが

for(min)
blueとred二つ積む
0個でない事に気をつけつつ、片方を積んだ場合をそれぞれ考える。

何でこれで良いかと言うと、二つ積んだ時は高さが一緒な為、同じブロックタワーとみなされる。
この考え方に赤が一番下の場合と、青が一番下の場合を区別しなくて済む。
Eulerdoraさんのコードを参考に直した。

import java.util.*;

public class TheBrickTowerEasyDivTwo {

	public int find(int redCount, int redHeight, int blueCount, int blueHeight) {
		int min = Math.min(blueCount, redCount);
		TreeSet<Integer> set = new TreeSet<Integer>();
		set.add(redHeight);
		set.add(blueHeight);

		for(int i = 0; i < min; i++){
			blueCount--;
			redCount--;
			int temp = (redHeight + blueHeight) *(i+1);
			set.add(temp);
			if(blueCount > 0) set.add(temp + blueHeight);
			if(redCount > 0) set.add(temp + redHeight);
		}
		return set.size();
	}

}