codingtest

11401번 - 이항 계수 3

2024-07-28
4분 분량
JAVA백준

문제 링크

문제 풀이 시간 : 2시간

문제 요약

  • 자연수 N과 정수 K가 주어졌을 때 이항계수를 1,000,000,007로 나눈 나머지를 구해라
  • 1 ≤ N ≤ 4,000,000
  • 0 ≤ K ≤ N

문제 풀이

이 문제는 알고리즘 시간때 배웠던 dp 알고리즘으로 풀면 되는 아주 간단한 문제 아닌가?

하고 신나게 바로 풀기 시작했다.

b[i][j] = b[i-1][j] + b[i-1][j-1]

초기 코드 (dp)

java
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int N = Integer.parseInt(st.nextToken());
        int K = Integer.parseInt(st.nextToken());

        int[][] bin = new int[N + 1][K + 1];
        for (int i = 0; i <= N; i++) {
            for (int j = 0; j <= Math.min(i, K); j++) {
                if (j == 0 || j == i) {
                    bin[i][j] = 1;
                } else {
                    bin[i][j] = (bin[i - 1][j] + bin[i - 1][j - 1]) % 1_000_000_007;
                }
            }
        }
        System.out.println(bin[N][K]);

    }
}

이렇게 아주 간단하게 코드를 작성할 수 있었고,

바로 메모리 초과를 받게 되었다.ㅋ

생각해보면 당연한게 N의 범위가 4,000,000이라서 배열이 최대 4,000,000 * 4,000,000이 될 수 있는 것이다.

그래서 생각한게 배열을 최소한으로 사용하면 메모리 초과를 해결할 수 있지 않을까?

K 길이의 배열을 두개만 사용해서 푸는 방법이다.

중간 코드(메모리 해결)

java
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {
    public static int N,K;

    public static void calc(int[] a, int[] b, int idx){
        for(int i=0;i<=Math.min(idx,K);i++){
            if(i==0||i==idx){
                a[i] = 1;
            }
            else{
                a[i] = b[i - 1] + b[i];
            }
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());

        int[] bin = new int[K + 1];
        int[] temp = new int[K + 1];

        for (int i = 0; i <= N; i++) {
            if(i%2==0){
                calc(bin,temp,i);
            }
            else{
                calc(temp, bin, i);
            }
        }

        int ans = (N % 2 == 0) ? bin[K] : temp[K];
        System.out.println(ans);
    }
}

이렇게 했더니 메모리 초과는 해결되었지만

시간초과가 뜬다.

그것도 그럴게 생각해보면 우리가 배운 dp 알고리즘은 O(n^2)의 시간복잡도를 가진다.

즉, 이 문제에서는 n 이 4,000,000이라서 엄청나게 느리게 동작한다는 것이다.

그러면 어떻게 해결해야 할까?

우리는 이항계수를 구하는 식을 알고 있다.

nCr=(nk)=n!r!(n−r)!nCr = {n \choose k} = \frac{n!}{r!(n-r)!}

위 식을 사용하면 해당 문제의 조건에서 O(n^2)의 시간복잡도를 거치지 않고 더 빠르게 구할 수 있다.

하지만 문제는 위 식에서 1,000,000,007로 나눈 나머지 값을 출력해야 한다.

우리는 보통 나머지 값을 출력할 때, 모듈러 연산을 사용한다.

모듈러 연산은 우리가 알고리즘 시간에 배운 것처럼 분배법칙이 성립한다.

하지만 모듈러 연산은 아래와 같이 나눗셈에서의 분배법칙은 성립되지 않는다.

(a÷b) mod m≠((a mod m)÷(b mod m)) mod m(a \div b)\bmod{m} \ne ((a\bmod{m})\div(b\bmod{m}))\bmod{m}

즉, 우리는 위의 나눗셈을 곱셉으로 바꿔줘야 한다.

어떻게 나눗셈을 곱셈으로 바꿀까?

역원을 이용하는 것이다.

ax=T→ax−1=T\frac{a}{x} = T \rightarrow ax^{-1}=T

위와 같이 x에 대한 역원을 구하면 손쉽게 나눗셈을 곱셈으로 표현해줄 수 있다.

즉, 우리가 구하고자 하는 식을 곱셈으로 바꾼다면

nCr=(nk)=n!r!(n−r)!=((n! mod p)×(k!(n−k)!)−1 mod p) mod pnCr = {n \choose k} = \frac{n!}{r!(n-r)!}=((n!\bmod{p})\times(k!(n-k)!)^{-1}\bmod{p})\bmod{p}

로 나타낼 수 있다.

그렇다면 우리가 이제 구해야 하는 것은 k!(n-k)!의 역원을 구하는 것이다.

여기서 우리는 ‘페르마의 소정리’를 사용하게 된다.

페르마의 소정리란?

a는정수,p는소수이며a∤p일때,ap−1≡1 mod pa는 정수, p는 소수이며 a\nmid p일 때,\\ a^{p-1} \equiv 1\bmod{p}

여기서 나뭇가지 같은 기호는 a가 p의 배수가 아닌 정수라는 뜻이다.

위 식의 예를 들어 p = 7, a = 12라고 한다면,

127−1=126=2985984≡1(mod7)12^{7-1} = 12^6 = 2985984 \equiv 1 \pmod{7}

이 된다.

즉, 12^6 을 7로 나눈 나머지가 1이 된다는 뜻이다.

그럼 이 식을 증명해보자.

다음과 같은 p-1 크기의 집합을 살펴보자. (단, p는 소수)

{1,2,⋯ ,p−1}\{1,2,\cdots,p-1\}

이 집합의 모든 원소에 a를 곱한다. (단, a는 p의 배수가 아닌 정수)

{a,2a,⋯ ,(p−1)a}\{a,2a,\cdots,(p-1)a\}

그럼 아래와 같은 식이 성립되게 되는데,

{1,2,⋯ ,p−1}={a,2a,⋯ ,(p−1)a}(modp)\{1,2,\cdots,p-1\} = \{a,2a,\cdots,(p-1)a\}\pmod{p}

예를 들어, p = 5, a = 2 라고 한다면,

{1,2,3,4}×2={2,4,6,8}{2,4,6,8} mod 5={2,4,1,3}={1,2,3,4}\{1,2,3,4\}\times2 = \{2,4,6,8\}\\ \{2,4,6,8\}\bmod{5} = \{2,4,1,3\} = \{1,2,3,4\}

가 된다.

이것을 증명하기 위해 아래 사실을 보인다.

{a,2a,⋯ ,(p−1)a}(modp)의모든원소들은서로다르다.\{a,2a,\cdots,(p-1)a\} \pmod{p}의 모든 원소들은 서로 다르다.

귀류법으로 해당 사실을 증명해보인다.

i≠jia=ja(modp)i \ne j\\ia=ja\pmod{p}

인 i,j 가 존재한다고 가정하자.

p가 소수이므로 a와 p는 서로소이다. 따라서 양변을 a로 나누면

i≡j(modp)i \equiv j\pmod{p}

가 된다.

여기서 i와 j는 같은 나머지를 가진다는 뜻이 된다.

그러나 1 ≤ i,j ≤ p-1 의 범위에서 i와 j의 값이 다르다는 것이 말이 안된다. 즉, i와 j는 동일하다가 된다.

따라서,

{a,2a,⋯ ,(p−1)a}{1,2,⋯ ,p−1}\{a,2a,\cdots,(p-1)a\}\\\{1,2,\cdots,p-1\}

위의 두 집합은 모든 원소들이 서로 다른 값을 가지는 같은 집합이라는 뜻이다.

따라서 두 집합의 모든 원소를 곱한 값도 같은 것이다.

ap−1(p−1)!≡(p−1)!(modp)a^{p-1}(p-1)! \equiv (p-1)!\pmod{p}

여기서 양변을 (p-1)!로 나누면 아래의 식이 성립하게 된다.

ap−1≡1(modp)a^{p-1} \equiv 1\pmod{p}

이게 페르마의 소정리이다.

사실 나도 잘 이해가 안된다.ㅎㅎ

풀이(계속)

위 정리를 사용해서 풀이를 계속해보자면-

ap−1≡1(modp)=a×ap−2≡1(modp)a^{p-1} \equiv 1\pmod{p}\\= a\times a^{p-2}\equiv 1 \pmod{p}

가 된다.

즉, a (mod p) 의 역원은 a^p-2 (mod p)가 되는 것이다.

그렇다면 우리가 구하고자 했던 역원은?

a=(k!(n−k)!),p=1,000,000,007(k!(n−k)!)−1=(k!(n−k)!)1000000007−2a = (k!(n-k)!), p = 1,000,000,007\\(k!(n-k)!)^{-1} = (k!(n-k)!)^{1000000007-2}

가 된다.

즉, 최종적으로 우리가 구해야하는 식은 아래와 같다.

n!k!(n−k)! mod p=(n!×(k!(n−k)!)−1) mod p=((n! mod p)×(k!(n−k)!) mod p) mod p=((n! mod p)×(k!(n−k)!)1000000007−2 mod p) mod p\begin{matrix} \frac{n!}{k!(n-k)!}\bmod{p} &=& (n! \times (k!(n-k)!)^{-1})\bmod{p}\\&=&((n!\bmod{p})\times(k!(n-k)!)\bmod{p})\bmod{p}\\&=&((n!\bmod{p})\times (k!(n-k)!)^{1000000007-2}\bmod{p})\bmod{p} \end{matrix}

이 된다.

이 식을 이용해서 코드를 짜면 아래와 같이 쉽게 구할 수 있게 된다.

최종 코드

java
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {
    public static int N, K;
    public static long DIV = 1_000_000_007;

    public static long pow(long a, long b) {
        if (b == 1) {
            return a % DIV;
        }
        long temp = pow(a, b / 2);
        if (b % 2 == 0) {
            return temp * temp % DIV;
        }
        return (temp * temp % DIV) * a % DIV;
    }

    public static long fact(int a) {
        long result = 1L;
        while (a > 1) {
            result = (result * a) % DIV;
            a -= 1;
        }
        return result;
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());
        K = Integer.parseInt(st.nextToken());

        long up = fact(N) % DIV;
        long down = fact(K) * fact(N - K) % DIV;

        System.out.println(up * pow(down, DIV - 2) % DIV);
    }
}

함께 읽으면 좋은 글

코딩 테스트2026-04-02

5430번 - AC

함수 R은 배열을 뒤집는 연산 함수 D는 첫 번째 원소를 버리는 연산 단, 배열이 비어 있을 때 D를 수행하면 error 주어진 함수 문자열을 순서대로 수행한 뒤 최종 배열을 출력하는 문제 처음에는 문제 설명 그대로 R이 나올 때마다 배열을 실제로 뒤집고, D가 나오면 맨 앞…

코딩 테스트2026-04-02

1966번 - 프린터 큐

여러 문서가 큐에 들어 있고, 각 문서마다 중요도가 있음 맨 앞 문서를 확인했을 때, 더 높은 중요도의 문서가 뒤에 하나라도 있으면 맨 뒤로 보냄 그렇지 않으면 바로 인쇄 특정 위치의 문서가 몇 번째로 인쇄되는지 구하는 문제 이 문제는 프린터의 동작을 그대로 구현하면 되는…

코딩 테스트2026-03-29

1158 - 요세푸스 문제

1번부터 N번까지 사람이 원형으로 앉아 있음 순서대로 K번째 사람을 제거 모든 사람이 제거될 때까지 반복 제거되는 순서를 출력하는 문제 이 문제는 원형으로 순회하면서 K번째를 제거하는 과정을 그대로 구현하면 된다.

코딩 테스트2026-03-29

2164번 - 카드2

1부터 N까지 카드가 순서대로 쌓여 있음 맨 위 카드를 버리고, 그다음 맨 위 카드를 맨 아래로 옮기는 과정을 반복 마지막에 남는 카드 번호를 구하는 문제 이 문제는 규칙을 찾기보다, 문제에서 하라는 과정을 그대로 구현하면 된다. 카드 더미에서 반복되는 동작은 두 가지다.

← 서재 목록으로 돌아가기