069.md

069 - Colorful Blocks 2 (★3)

解答 1

冪剰余 (base^exponent mod m) は、繰り返し二乗法と呼ばれる手法を用いると、巨大な整数を登場させずに高速に計算できる。

#include <iostream>

// 冪剰余 (base^exponent mod m) を計算 (繰り返し二乗法)
int ModularPow(long long base, unsigned long long exponent, int m)
{
	if (m == 1)
	{
		return 0;
	}

	int result = 1;
	base %= m;

	while (exponent)
	{
		if (exponent & 1)
		{
			result = (result * base) % m;
		}

		exponent >>= 1;
		base = (base * base) % m;
	}

	return result;
}

int main()
{
	// N 個のブロック, K 種類の色
	long long N, K;
	std::cin >> N >> K;

	constexpr int MOD = 1'000'000'007;

	if (K == 1) // コーナーケース: 色が 1 種類
	{
		if (N == 1) // ブロックが 1 個
		{
			std::cout << 1 << '\n';
		}
		else // ブロックが複数
		{
			std::cout << 0 << '\n'; // 塗れない
		}
	}
	else if (N == 1)
	{
		// K 通り
		std::cout << K << '\n';
	}
	else if (N == 2)
	{
		// K * (K-1) 通り
		std::cout << (K * (K - 1) % MOD) << '\n';
	}
	else
	{
		// K * (K-1) * (K-2)^(N-2) 通り
		std::cout << (K * (K - 1) % MOD
			* ModularPow(K - 2, N - 2, MOD) % MOD) << '\n';
	}
}

解答 2 (AC Library を使用)

解答 1 と同様に繰り返し二乗法で冪剰余を計算する関数が、AC Library に atcoder::pow_mod(base, exponent, m) として実装されています。

#include <iostream>
#include <atcoder/math.hpp> // atcoder::pow_mod()

int main()
{
	// N 個のブロック, K 種類の色
	long long N, K;
	std::cin >> N >> K;

	constexpr int MOD = 1'000'000'007;

	if (K == 1) // コーナーケース: 色が 1 種類
	{
		if (N == 1) // ブロックが 1 個
		{
			std::cout << 1 << '\n';
		}
		else // ブロックが複数
		{
			std::cout << 0 << '\n'; // 塗れない
		}
	}
	else if (N == 1)
	{
		// K 通り
		std::cout << K << '\n';
	}
	else if (N == 2)
	{
		// K * (K-1) 通り
		std::cout << (K * (K - 1) % MOD) << '\n';
	}
	else
	{
		// K * (K-1) * (K-2)^(N-2) 通り
		std::cout << (K * (K - 1) % MOD
			* atcoder::pow_mod(K - 2, N - 2, MOD) % MOD) << '\n';
	}
}