Java 문제풀이 - 나머지 합

문제

수 N개 A1, A2, ..., AN이 주어진다. 이때, 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 구하는 프로그램을 작성하시오.

즉, Ai + ... + Aj (i ≤ j) 의 합이 M으로 나누어 떨어지는 (i, j) 쌍의 개수를 구해야 한다.

입력

첫째 줄에 N과 M이 주어진다. (1 ≤ N ≤ 106, 2 ≤ M ≤ 103)

둘째 줄에 N개의 수 A1, A2, ..., AN이 주어진다. (0 ≤ Ai ≤ 109)

출력

첫째 줄에 연속된 부분 구간의 합이 M으로 나누어 떨어지는 구간의 개수를 출력한다.

예제 입/출력

나의 풀이

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

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer stringTokenizer = new StringTokenizer(bufferedReader.readLine());

        int N = Integer.parseInt(stringTokenizer.nextToken());
        int M = Integer.parseInt(stringTokenizer.nextToken());

        long[] sumArray = new long[N + 1];
        long[] remainderArray = new long[M];
        long numOfDividedByM = 0;

        stringTokenizer = new StringTokenizer(bufferedReader.readLine());
        for (int i = 1; i < N + 1; i++) {
            sumArray[i] = sumArray[i - 1] + Integer.parseInt(stringTokenizer.nextToken());
            remainderArray[(int)(sumArray[i] % M)]++;
        }

        // 나머지가 0인 경우
        long remainderZeroNum = remainderArray[0];
        numOfDividedByM += remainderZeroNum * (remainderZeroNum - 1) / 2 + remainderZeroNum;

        // 나머지가 0이 아닌 경우
        for (int i = 1; i < M; i++) {
            long remainderNum = remainderArray[i];
            numOfDividedByM += remainderNum * (remainderNum - 1) / 2;
        }

        System.out.println(numOfDividedByM);
    }
}

예시 풀이

import java.io.IOException;
import java.util.Scanner;
public class P10986_나머지합 {
  public static void main(String[] args) throws IOException {
    Scanner sc = new Scanner(System.in);
    int N = sc.nextInt();
    int M = sc.nextInt();
    long[] S = new long[N];
    long[] C = new long[M];
    long answer = 0;
    S[0] = sc.nextInt();
    for (int i = 1; i < N; i++) { // 수열 합배열 만들기
      S[i] = S[i - 1] + sc.nextInt();
    }
    for (int i = 0; i < N; i++) { // 합배열의 모든 값에 %연산 수행하기
      int remainder = (int) (S[i] % M);
      if (remainder == 0)
        answer++; // 0~i까지의 구간합 자체가 0인 경우 정답에 더해주기
      C[remainder]++; // 같은 나머지를 가진 인덱스의 개수 카운팅 해주기
    }
    for (int i = 0; i < M; i++) {
      if (C[i] > 1) {
        answer = answer + (C[i] * (C[i] - 1) / 2); // 같은 나머지를 가진 인덱스들중 2개를 뽑는 경우의 수를 더해주기
      }
    }
    System.out.println(answer);
  }
}

리뷰

M으로 나눴을 때 나머지가 같은 두 값을 빼면 M으로 정확히 나눠진다는 발상이 중요한 문제였다.

사실 이 책이 수록된 문제집에서 조건에 오타가 있어서 처음에 조금 헤맸다. 백준 온라인 문제 설명 위주로 보는게 맞는 것 같다. 그리고 이 참에 책에 정오표를 전부 적용해서 수정해야겠다.

이 문제를 풀면서 몇 가지 습관을 들여야겠다고 느낀 것이 있다.
1. 문제를 읽고 자료형을 int가 아닌 long으로 선언하자.
2. 스스로 코드를 짜더라도 항상 주석을 같이 써보자.
3. 무작정 코드부터 짜지 말고, 문제를 좀만 더 꼼꼼하게 읽고 생각해보자.