[HackerRank] The Power Sum (JAVA)

문제

Find the number of ways that a given integer, X, can be expressed as the sum of the N^th powers of unique, natural numbers. For example, if X = 13 and N = 2, we have to find all combinations of unique squares adding up to 13. The only solution is 2² + 3².

 

 

입력

The first line contains an integer X. The second line contains an integer N.

 

 

제한

1 <= X <= 1000

 

2 <= N <= 10

 

 

풀이

조합 혹은 다이나믹 프로그래밍을 이용하여 풀 수 있는 문제였습니다.

 

 

조합을 이용한 풀이

문제에서 N은 지수를 뜻하는데, 어떠한 수들의 N제곱을 한 값을 모두 더해서 X를 만드는 방법이 몇 가지인지 찾아야합니다. 그래서 가장 처음 생각할 수 있는 방법은 무식하게 다 해보는 것입니다.

 

우선, 최악의 경우를 먼저 생각해 보자면, X = 1000이고 N = 2일 때입니다. 그리고 무식하게 다 해 본다고 가정하면, ₁₀₀₀C₁ 에서 ₁₀₀₀C₁₀₀₀까지 조합을 이용하여 1부터 1000까지의 수 중 임의의 수들을 제곱하여 더해 보면서 X와 같은지 비교할 것입니다.

 

 

하지만, 우리는 여기서 연산 횟수를 줄일 수 있습니다. 첫 번째로, 1에서 X까지가 아니라 1에서 X의 N제곱근까지로 줄일 수 있습니다. 가령, X = 100이고 N = 3이라면 1에서 4까지의 수에서만 몇 가지 뽑아서 연산해 보면 됩니다.

 

두 번째로, 조합에서 선택될 대상의 개수를 줄일 수 있습니다. 위의 예에서 1부터 1000까지 수들을 모조리 뽑아보면서 연산을 한다고 하였는데, 이렇게 하면 수행 시간이 어마무시하게 길어집니다. 이때, 1^N + 2^N + ... + k^N을 더해서 이 합이 X보다 작되, 최댓값을 만드는 k를 구하면 됩니다. 가령, X = 100이고 N = 3이라면 k는 6입니다.

 

왜 1부터 k까지 N제곱하여 더한 값을 X와 비교하는 걸까요? 그것은 선택될 대상의 최대 개수를 구하기 위함입니다.

 

 

이제, 위 두 가지 최적화 방법을 통한 조합을 이용하여 X를 만들 수 있는 방법을 구하시면 됩니다. 참고로, 저는 백트래킹으로 조합을 구현하였고, 자세한 코드는 아래를 참조하시면 되겠습니다.

 

 

조합을 이용한 소스코드

	import java.io.*;
	import java.math.*;
	import java.security.*;
	import java.text.*;
	import java.util.*;
	import java.util.concurrent.*;
	import java.util.regex.*;
	public class Solution {
	static int X;
	static int N;
	static int maxNum;
	static int possibleCnt;
	static int ans;
	static void combination(int[] arr, boolean[] visited, int start, int n, int r) {
	if (r == 0) {
	if (powerSum(arr, visited)) {
	ans++;
	}
	return;
	}
	for (int i = start; i <= n; i++) {
	visited[i] = true;
	combination(arr, visited, i + 1, n, r - 1);
	visited[i] = false;
	}
	}
	static boolean powerSum(int[] arr, boolean[] visited) {
	int total = 0;
	for (int i = 1; i <= maxNum; i++) {
	if (visited[i]) {
	total += Math.pow(arr[i], N);
	}
	if (total > X) {
	return false;
	}
	}
	return total == X;
	}
	public static void main(String[] args) throws IOException {
	BufferedWriter bw = new BufferedWriter(new FileWriter(System.getenv("OUTPUT_PATH")));
	BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
	X = Integer.parseInt(br.readLine());
	N = Integer.parseInt(br.readLine());
	maxNum = (int) Math.pow(X, (1 / (double) N));
	possibleCnt = 0;
	for (int i = 1, total = 0; ; i++, possibleCnt++) {
	total += Math.pow(i, N);
	if (total > X) {
	break;
	}
	}
	int[] arr = new int[maxNum + 1];
	for (int i = 1; i <= maxNum; i++) {
	arr[i] = i;
	}
	ans = 0;
	for (int i = 1; i <= possibleCnt; i++) {
	boolean[] visited = new boolean[maxNum + 1];
	combination(arr, visited, 1, maxNum, i);
	}
	bw.write(ans + "\n");
	bw.flush();
	bw.close();
	br.close();
	}
	}

view raw The Power Sum - comb.java hosted with ❤ by GitHub

 

 

 

다이나믹 프로그래밍을 이용한 풀이

조합은 최악의 경우 약 9억 번의 연산이 필요하고, 대략 1억 번의 연산당 1초가 걸리므로 9초정도 시간이 소요될 수 있습니다. 저는 사실 반신반의하면서 제출하였는데 통과를 하였고, 좀 더 효율적인 방법을 생각해 보았습니다.

 

그리고 bottom-up 방식을 이용한 다이나믹 프로그래밍으로 풀 수 있다는 사실을 알게 되었습니다. 로직 자체는 단순합니다. dp[i] = n으로 i번째 수를 만드는 n가지 방법으로 정의하면 됩니다.

 

처음에 1부터 X의 N제곱근까지의 값들에 대해서 N제곱을 취한뒤, 리스트에 넣어 둡니다. X = 100이고, N = 3일 경우 리스트에는 {1, 2, 3, 4}가 있을 것입니다.

 

그리고 dp[0] = 1, dp[1] = ... = dp[n] = 0으로 초기화한 뒤, 리스트의 요소를 처음부터 탐색합니다. 리스트의 요소를 c라 하고, dp[X - c] 값이 0이 아닌 경우를 찾습니다. 이전에 만든 수에서 특정 수를 더하면 X를 만들 수 있는지 bottom-up 방식으로 다이나믹 프로그래밍을 해 나가는 것입니다. 수식으로는 dp[X] += dp[X - c]가 되겠군요.

 

 

아래는 위 과정을 정리한 소스코드입니다.

 

 

다이나믹 프로그래밍을 이용한 소스코드

	import java.io.*;
	import java.math.*;
	import java.security.*;
	import java.text.*;
	import java.util.*;
	import java.util.concurrent.*;
	import java.util.regex.*;
	public class Solution {
	// Complete the powerSum function below.
	static int powerSum(int X, int N) {
	int maxNum = (int) Math.pow(X, 1 / (double) N);
	ArrayList<Integer> arrList = new ArrayList<>();
	for (int i = 1; i <= maxNum; i++) {
	arrList.add((int) Math.pow(i, N));
	}
	int[] dp = new int[X + 1];
	dp[0] = 1;
	for (int i = 0; i < arrList.size(); i++) {
	int now = arrList.get(i);
	for (int j = X - now; j >= 0; j--) {
	if (dp[j] != 0) {
	dp[j + now] += dp[j];
	}
	}
	}
	return dp[X];
	}
	private static final Scanner scanner = new Scanner(System.in);
	public static void main(String[] args) throws IOException {
	BufferedWriter bufferedWriter = new BufferedWriter(new FileWriter(System.getenv("OUTPUT_PATH")));
	int X = scanner.nextInt();
	scanner.skip("(\r\n|[\n\r\u2028\u2029\u0085])?");
	int N = scanner.nextInt();
	scanner.skip("(\r\n|[\n\r\u2028\u2029\u0085])?");
	int result = powerSum(X, N);
	bufferedWriter.write(String.valueOf(result));
	bufferedWriter.newLine();
	bufferedWriter.close();
	scanner.close();
	}
	}

view raw The Power Sum - dp.java hosted with ❤ by GitHub

 

 

지적 혹은 조언 환영합니다! 언제든지 댓글로 남겨주세요.