백준 14465 - 소가 길을 건너간 이유 5

int idx = 0;
sort(lights.begin(), lights.end());
for (int i = 1; i <= N; i++)
{
	if (lights[idx] == i)
	{
		sum[i] = sum[i - 1];
		idx++;
	}
	else 
	{
		sum[i] = sum[i - 1] + 1;
	}
}​

문제 설명

농부 존의 농장에 원형 길이 있다고 했지만, 길은 그뿐만이 아니다. 그 옆에 일자형 길이 있는데, 1번부터 N번까지의 번호가 붙은 횡단보도 N (1 ≤ N ≤ 100,000)개로 이루어져 있다. 교통사고를 방지하기 위해 존은 각 횡단보도에 신호등을 설치해 놓았다. 그러던 어느 날, 강력한 뇌우로 인해 몇몇 신호등이 망가졌다. 존은 연속한 K개의 신호등이 존재하도록 신호등을 수리하고 싶다. 이번에도 우리가 존을 도와주자.

https://www.acmicpc.net/problem/14465

---

 

 

 

 

 

 

제한 사항

---

 

 

 

 

 

 

풀이

문제를 요약하면 고장난 신호등의 번호가 주어졌을 때 x개를 수리하여 연속된 K개의 정상 신호등을 만드려 할 때 x의 최솟값을 구하면 된다.

 

일단 시간을 생각하지 않고 문제를 풀었을 때 고장난 신호등을 정렬한 후 K 만큼씩 확인하며 최솟값을 구하면 된다.

이 경우 $O(N * K)$의 시간이 걸린다.

하지만 문제에 N에 대한 제한이 없고 아마 시간 초과가 발생하도록 입력이 주어질 것 같다.

 

따라서 시간을 줄이기 위해 다른 방법을 사용해 보자.

내가 선택한 방법은 누적합을 이용하는 것이다.

정상인 신호등의 개수를 누적합을 통해 계산한다면 x-k~x구간의 멀쩡한 신호등의 개수를 $O(1)$만에 구할 수 있다.

누적합을 구하는 과정까지 더한다면 $O(N)$만에 정답을 구할 수 있는 것이다.

 

우선 누적합을 구하는 과정은 다음과 같다.

int idx = 0;
sort(lights.begin(), lights.end());
for (int i = 1; i <= N; i++)
{
	if (lights[idx] == i)
	{
		sum[i] = sum[i - 1];
		idx++;
	}
	else 
	{
		sum[i] = sum[i - 1] + 1;
	}
}

 

이를 이용하여 N개의 신호등을 살펴보며 K개의 범위에 정상적인 신호등의 개수를 세어 K에서 빼주면 된다.

int ans = MAX;
for (int i = K; i <= N; i++)
{
	ans = min(ans, K - (sum[i] - sum[i - K]));
}

---

 

 

 

 

 

전체 코드

#include <bits/stdc++.h>
#include <unordered_set>
using namespace std;
#define INPUT_OPTIMIZE cin.tie(NULL); cout.tie(NULL); ios::sync_with_stdio(false);
#define INF 2e9
#define MAX 987654321

using namespace std;
int N, K, B;
vector<int> lights;
vector<int> sum;

int main()
{
	INPUT_OPTIMIZE;
	
	cin >> N >> K >> B;

	lights.resize(B);
	sum.resize(N + 1);

	for (int i = 0; i < B; i++)
	{
		cin >> lights[i];
	}

	int idx = 0;
	sort(lights.begin(), lights.end());
	for (int i = 1; i <= N; i++)
	{
		if (lights[idx] == i)
		{
			sum[i] = sum[i - 1];
			idx++;
		}
		else 
		{
			sum[i] = sum[i - 1] + 1;
		}
	}

	int ans = MAX;
	for (int i = K; i <= N; i++)
	{
		ans = min(ans, K - (sum[i] - sum[i - K]));
	}

	cout << ans;
	
	return 0;
}

---