백준 1241 머리 톡톡 python, java

N이 최대 10만.. 부르트포스로는 풀 수가 없다.
그러가다 생각난게 아리토스테네스의 체이다.

📜 문제 바로 가기 : 머리 톡톡

제출코드

파이썬

import sys
input = sys.stdin.readline

from collections import defaultdict

N = int(input())
arr = [int(input()) for _ in range(N)]

maxValue = max(arr)
cnt_arr = [0] * (maxValue+1)
cnt_dict = defaultdict(int)

for i in arr:
    cnt_dict[i] += 1

for i in set(arr):
    k = 2*i
    while k <= maxValue:
        cnt_arr[k] += cnt_dict[i]
        k += i

for i in arr:
    print(cnt_arr[i] + cnt_dict[i]-1)

자바

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

public class Main {


    public static void main(String[] args) throws IOException{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        int N = Integer.parseInt(br.readLine());
        int maxValue = -1;
        HashMap<Integer, Integer> visitMap = new HashMap<>();

        int[] arr = new int[N];
        for (int i = 0; i < N; i++) {
            int number = Integer.parseInt(br.readLine());
            arr[i] = number;
            if (number > maxValue) {
                maxValue = number;
            }

            visitMap.put(number, 1);
        }

        int[] cnt_arr = new int[maxValue+1];
        HashMap<Integer, Integer> cnt_map = new HashMap();

        for (int num : arr) {
            if (cnt_map.containsKey(num) == true) {
                int tmp_key = cnt_map.get(num);
                cnt_map.put(num, tmp_key+1);
            } else {
                cnt_map.put(num, 1);
            }
        }

        for (int num : arr) {
            if (visitMap.get(num) == 1) {
                visitMap.put(num, 0);
                int k = num + num;
                while (k <= maxValue) {
                    cnt_arr[k] += cnt_map.get(num);
                    k += num;
                }
            }
        }

        for (int num : arr) {
            System.out.println(cnt_arr[num] + cnt_map.get(num) - 1);
        }
    }
}

---

접근방법

나를 제외한 다른 사람의 숫자를 보고 내 숫자가 그 숫자의 배수이면 머리를 친다. (....)
그리고 내가 톡톡한 사람의 숫자를 카운트해서 출력한다.

하나하나 다 확인하게 된다면 O(N^2) 만큼의 시간이 필요하고 N은 10만이기 때문에 이 방법은 안된다.

그러다가 문득 "아리토스테네스의 체"가 생각났다. 이 알고리즘은 소수를 찾는 알고리즘인데 가장 작은 소수인 2부터 올라가면서 n이 소수이면 2n, 3n ... 은 소수가 안되도록 만드는 방법을 사용한다.
이 방법처럼 내 숫자를 보고 내 머리를 칠 수 있는 번호에 카운트를 더해주는 방법이다. 즉, 내 배수의 숫자에 모두 +1 을 해준다.

이방법에서 주의할 점이 있는데 중복된 숫자를 가진 경우이다. 예를 들어 최대숫자가 100만인데 1을 적은 사람이 100명만 있어도 연산이 1억이된다. 이러한 문제를 피하기 위해 중복숫자를 카운트할 수 있는 딕셔너리를 만들었다. (자바에선 HashMap)
이를 통해 내 머리를 칠 수 있는 사람에게 +1을 해줄 때 +1이 아닌 나와 같은 숫자를 적은 사람만큼 더해주면서 불필요한 연산을 줄였다.
또한 x는 x의 배수이기 때문에 x번호를 가진 사람의 톡톡 횟수는 x를 가진 사람의 수 -1 만큼 추가로 더해주는것을 잊지 말아야 한다.