코딩테스트 스터디 - 백준 25406번 : 식사 계획 세우기 (C++)

문제

링크 : https://www.acmicpc.net/problem/25406

25406번: 식사 계획 세우기

요약

식당 N개가 각자 1~N 음식 중 하나를 팔고 있다.

모든 식당을 다 방문하면서 같은 음식을 두 번 연속 먹지 않는 식사계획을 세우려한다.

이런 식사계획 중 사전순으로 가장 빠른 식사계획을 출력해야한다.

의식의 흐름

결론부터 말하자면 혼자 힘으로 풀지 못했다.

 

우선 올바른 식사계획을 세우기 위한 조건은 '어떤 음식의 식당 수가 과반이 넘지 않는 것'이다.

그렇기 때문에 기본적으로는 식당 번호가 가장 작으면서 이전 음식이랑 안겹치는 식당들을 고르다가,

어느 순간 어떤 음식의 식당 수가 과반이 되면 그 식당을 고르면 된다.

 

이걸 구현하기 위해선 음식별로 식당 list를 가지는 dictionary of list가 필요하다.

하지만 이렇게 구현하면 한 번 식당을 고를 때,

가장 많이 남은 식당 수 고르기 (N) * 식당 번호가 가장 작은 식당 고르기 (N) 이 돼 총 O(N^3)이 된다.

그래서 계속 17점을 받았고, 여러군데에 캐싱을 하는 시도를 해봤지만 결과는 같았다.

 

결국 검색찬스를 써서 C++ set를 이용한 구현 방법을 배웠다.

참고 블로그 : https://velog.io/@ks1ksi/백준-25406번-식사-계획-세우기

Python에는 binary search tree를 구현한 자료구조가 없어서, 그걸 만드느니 그냥 C++로 풀기로 했다.

 

핵심은 최대와 최소를 계속 구해야하는 식당 수와 식당 index를 set에 넣어 이 연산을 logN 속도로 만드는 것이다.

그러면 총 O(NlogN)으로 풀 수 있다.

구현

자료구조

shops : 음식 별 남은 식당을 보관하는 map of queues. 식당 입력이 번호순으로 되기 때문에 enqueue만으로 정렬이 되고, 음식 x에 대해 가장 빠른 식당은 shops[x].front()로 알 수 있다.

counts : (남은 식당 수, 음식 종류) 를 보관하는 set

idxs : (이 음식의 첫 식당, 음식 종류) 를 보관하는 set

 

음식 종류에서 식당이 하나 빠질 때마다, 대응하는 원소를 counts와 idx에서 업데이트시켜줘야한다.

가장 많이 남은 식당은 counts.rbegin()으로, 가장 번호가 빠른 식당은 idxs.begin()으로 구할 수 있어 상수 속도가 된다.

이후 대응하는 원소를 counts와 idx에서 찾는 건 logN의 시간이 걸린다.

알고리즘

우선 맨 처음에 올바른 식사계획을 만들 수 있는지 체크해 -1을 출력한다.

이후에는 만들 수 있다는 가정하에, 식당을 고르는 작업을 N번 반복한다.

 

1. counts에서 제일 큰 원소의 count값이 과반일 경우

이 음식을 파는 가장 숫자가 작은 식당을 고른다.

2. 아닐 경우

idxs를 순회해 마지막으로 고른 음식과 다른 음식이 나오면, 그 식당을 고른다.

 

식당을 골랐으면 그에 맞는 값들을 counts와 idxs에서 삭제해준다.

그 다음 shops에서 그 식당을 빼주고, 

남은 식당 수와 가장 빠른 식당을 update해 다시 counts와 idxs에 insert 해준다. (남은 식당이 0이면 안넣는다.)

코드

이런 구현 은 간단하지만 시간이 빡빡한 문제를 두 번째로 풀어보는데,

진짜 문제 요구조건에 딱 맞는 자료구조를 선택하고, 캐싱을 활용해야 풀 수 있는 것 같다.

 

/*
	문제 출처 : https://www.acmicpc.net/problem/25406
*/

#define _CRT_SECURE_NO_WARNINGS
#include <stdio.h>
#include <set>
#include <queue>
#include <map>

int main()
{
	int N;
	scanf("%d", &N);

	std::map<int, std::queue<int>> shops;

	for (int i = 0; i < N; i++)
	{
		int food;
		scanf("%d", &food);
		if (shops.find(food) == shops.end())
			shops[food] = std::queue<int>();
		shops[food].push(i + 1);
	}

	std::set<std::pair<int, int>> counts;	//남은 수, 음식
	std::set<std::pair<int, int>> idxs;		//첫 식당, 음식

	for (int i = 1; i <= N; i++)
	{
		if (shops.find(i) == shops.end())
			continue;
		idxs.insert({ shops[i].front(),i });
		counts.insert({ shops[i].size(),i });
	}

	if ((*(counts.rbegin())).first > (int)((N + 1) / 2))
	{
		printf("-1");
		return 0;
	}
	int lastFood = -1;
	for (int i = 0; i < N; i++)
	{
		auto item = counts.rbegin();
		int count, food, idx;
		//과반이 있음
		if (item->first > (int)((N - i) / 2))
		{
			count = item->first;
			food = item->second;
			idx = shops[food].front();
		}
		else
		{
			auto it = idxs.begin();
			for (; it != idxs.end(); it++)
			{
				if (it->second != lastFood)
					break;
			}
			idx = it->first;
			food = it->second;
			count = shops[food].size();
		}

		counts.erase({ count, food });
		idxs.erase({ idx, food });

		printf("%d ", idx);
		lastFood = food;

		shops[food].pop();

		if (count - 1 > 0)
		{
			counts.insert({ count - 1, food });
			idxs.insert({ shops[food].front(), food });
		}
	}
	return 0;
}