TopCoder SRM 504.5 Div1 Medium: TheJackpotDivOne

解法

問題文通りに分配していけばいいが、全員同じになったら後は等分配するだけになるので、ループする必要はない。

コード

import java.math.BigInteger;
import java.util.Arrays;
import java.util.PriorityQueue;

public class TheJackpotDivOne {

	public long[] find(long[] money, long jackpot) {
		int N = money.length;
		if (N == 1) {
			money[0] += jackpot;
			return money;
		}

		PriorityQueue<BigInteger> priorityQueue = new PriorityQueue<>();
		BigInteger sum = BigInteger.valueOf(0);
		BigInteger max = BigInteger.valueOf(0);
		for (int i = 0; i < N; i++) {
			BigInteger friend = BigInteger.valueOf(money[i]);
			sum = sum.add(friend);
			priorityQueue.add(friend);
			max = max.max(friend);
		}

		while (jackpot > 0) {
			BigInteger poorest = priorityQueue.poll();

			// 全員同じになったら後は等分配するだけになるので終了する
			if (max.subtract(poorest).longValue() == 0) {
				priorityQueue.add(poorest);
				break;
			}

			BigInteger x = sum.add(BigInteger.valueOf(N));// sum+N
			x = x.divide(BigInteger.valueOf(N));// x/N
			x = x.subtract(poorest);// x-poorest
			x = x.min(BigInteger.valueOf(jackpot));
			jackpot -= x.longValue();
			priorityQueue.add(poorest.add(x));
			sum = sum.add(x);
			max = max.max(poorest.add(x));

		}
		for (int i = 0; i < N; i++) {
			BigInteger poorest = priorityQueue.poll();
			money[i] = poorest.longValue();
		}

		// 等分配するだけ
		if (jackpot > N) {
			long x = jackpot / N;
			for (int i = 0; i < N; i++) {
				money[i] += x;
			}
			jackpot -= x * N;
		}
		Arrays.sort(money);
		for (int i = 0; i < jackpot; i++) {
			money[i]++;
		}
		Arrays.sort(money);

		return money;
	}
}