백준 15961번 - 회전 초밥 [JAVA]

문제 요약

1. 벨트의 임의의 한 위치부터 k개의 접시를 연속해서 먹는다.
2. 쿠폰에 적혀진 초밥을 무료로 하나 제공한다.
3. 1번과 2번을 수행한 모든 경우의 수에 대해 가장 많이 먹을 수 있는 초밥 종류의 수를 구한다.

• 회전 초밥 벨트에 놓인 접시수: 8
• 초밥의 가짓수: 30
• 연속해서 먹는 접시의 수: 4
• 쿠폰 번호: 30

배열을 사용하고, 모듈로를 이용하여 회전하는 벨트 느낌으로 만들 수 있다. (원형 배열)

Naive

• 이중 포문을 사용해서 $k$개를 연속으로 선택하고, 초밥 가짓수를 세는 방법을 먼저 생각해볼 수 있다.

int n = in.nextInt(), m = in.nextInt(), k = in.nextInt(), c = in.nextInt();
int arr[] = new int[n];

for (int i = 0; i < n; i++) arr[i] = in.nextInt();
int ans = 0;
for (int i = 0; i < n; i++) { // 벨트 위치 지정
	Set<Integer> set = new HashSet<>();
	for (int j = i; j < i + k; j++) { // 연속한 k개의 초밥 먹기
		set.add(arr[j % n]);
	}
	set.add(c); // 쿠폰으로 무료 초밥 먹기
	ans = Math.max(ans, set.size());
}
out.print(ans);

근데 이렇게되면 시간복잡도는 $O(N*K)$가 된다.

 

관찰

• 4개중에 3개는 이전과 동일하다.
• 1개의 초밥이 추가되고, 1개의 초밥이 삭제된다.
• 4개가 아닌 $k$개로 일반화해도 동일하다.

여기서 드는 생각

• Naive하게 짰을 때, 초밥 종류 수를 Set으로 관리했는데, 추가할 때는 문제가 없지만 뺄 때 문제가 생긴다.
• 해당 구간에 같은 초밥 종류가 몇 개가 있는지 몰라서 생기는 문제다.
• 그러면 구간에 대한 각 초밥에 대한 개수를 셀 필요가 있겠군!
• $cnt$ 배열을 만들어서 관리해보자.
• 다음 초밥을 추가했을 때 (1 증가), $cnt$ 값이 1이면 새롭게 등장한 초밥이라는 거고,
• 이전 구간의 첫 초밥을 뺏을 때(1 감소), $cnt$ 값이 0이면 지금 구간에 더이상 해당 종류의 초밥이 없다는 뜻이다.

 

구현

• 회전하는 벨트이니까 배열이 계속 반복된다고 볼 수 있다. (모듈로 처리를 하면됨)
• 맨 왼쪽부터 하나씩 초밥을 추가해나가고
  • 추가했는데 $cnt$ 값이 1이면 초밥 가짓수 1 증가
• 초밥을 빼는 경우는 인덱스가 $k$ 일 때부터!
  • 뺏는데 $cnt$ 값이 0이면 초밥 가짓수 1 감소
• 우리가 실질적으로 체크해야 하는 배열의 인덱스는 $0$부터 $n + k - 2$까지

 

핵심 코드

int n = in.nextInt(), m = in.nextInt(), k = in.nextInt(), c = in.nextInt();
int arr[] = new int[n], cnt[] = new int[m + 1];
for (int i = 0; i < n; i++) arr[i] = in.nextInt();

int val = 0, ans = 0;
for (int i = 0; i <= n + k - 2; i++) {
    if (++cnt[arr[i % n]] == 1) val++; // 회전하니까 모듈로 연산
    if (i >= k && --cnt[arr[i - k]] == 0) val--;
    ans = Math.max(ans, val + (cnt[c] == 0 ? 1 : 0));
}
out.print(ans);

한 번의 탐색으로 정답을 구할 수 있다.

시간복잡도: $O(N + K)$