[프로그래머스] 두 큐 합 같게 만들기 - 파이썬(Python)

문제 설명

 

 

길이가 같은 두 개의 큐가 주어집니다. 하나의 큐를 골라 원소를 추출(pop)하고, 추출된 원소를 다른 큐에 집어넣는(insert) 작업을 통해 각 큐의 원소 합이 같도록 만들려고 합니다. 이때 필요한 작업의 최소 횟수를 구하고자 합니다. 한 번의 pop과 한 번의 insert를 합쳐서 작업을 1회 수행한 것으로 간주합니다.

큐는 먼저 집어넣은 원소가 먼저 나오는 구조입니다. 이 문제에서는 큐를 배열로 표현하며, 원소가 배열 앞쪽에 있을수록 먼저 집어넣은 원소임을 의미합니다. 즉, pop을 하면 배열의 첫 번째 원소가 추출되며, insert를 하면 배열의 끝에 원소가 추가됩니다. 예를 들어 큐 [1, 2, 3, 4]가 주어졌을 때, pop을 하면 맨 앞에 있는 원소 1이 추출되어 [2, 3, 4]가 되며, 이어서 5를 insert하면 [2, 3, 4, 5]가 됩니다.

다음은 두 큐를 나타내는 예시입니다.
queue1 = [3, 2, 7, 2]   |   queue2 = [4, 6, 5, 1]

두 큐에 담긴 모든 원소의 합은 30입니다. 따라서, 각 큐의 합을 15로 만들어야 합니다. 예를 들어, 다음과 같이 2가지 방법이 있습니다.

i) queue2의 4, 6, 5를 순서대로 추출하여 queue1에 추가한 뒤, queue1의 3, 2, 7, 2를 순서대로 추출하여 queue2에 추가합니다. 그 결과 queue1은 [4, 6, 5], queue2는 [1, 3, 2, 7, 2]가 되며, 각 큐의 원소 합은 15로 같습니다. 이 방법은 작업을 7번 수행합니다.
ii) queue1에서 3을 추출하여 queue2에 추가합니다. 그리고 queue2에서 4를 추출하여 queue1에 추가합니다. 그 결과 queue1은 [2, 7, 2, 4], queue2는 [6, 5, 1, 3]가 되며, 각 큐의 원소 합은 15로 같습니다. 이 방법은 작업을 2번만 수행하며, 이보다 적은 횟수로 목표를 달성할 수 없습니다.

따라서 각 큐의 원소 합을 같게 만들기 위해 필요한 작업의 최소 횟수는 2입니다.

길이가 같은 두 개의 큐를 나타내는 정수 배열 queue1, queue2가 매개변수로 주어집니다. 각 큐의 원소 합을 같게 만들기 위해 필요한 작업의 최소 횟수를 return 하도록 solution 함수를 완성해주세요. 단, 어떤 방법으로도 각 큐의 원소 합을 같게 만들 수 없는 경우, -1을 return 해주세요.

 

 

---

 

제한사항

 

 

• 1 ≤ queue1의 길이 = queue2의 길이 ≤ 300,000
• 1 ≤ queue1의 원소, queue2의 원소 ≤ 109
• 주의: 언어에 따라 합 계산 과정 중 산술 오버플로우 발생 가능성이 있으므로 long type 고려가 필요합니다.

 

 

---

 

입출력 예

 

 

queue1	queue2	result
[3, 2, 7, 2]	[4, 6, 5, 1]	2
[1, 2, 1, 2]	[1, 10, 1, 2]	7
[1, 1]	[1, 5]	-1

 

 

---

 

코드 설명

 

 

해당 문제를 풀이할 때 가장 먼저 고민한 점은 '두 큐의 합이 같아지게 만드는 작업의 최소 횟수를 어떻게 구할 수 있을까?' 였습니다.

 

문제의 ii)를 살펴보면 queue1에는 [3, 2, 7, 2], queue2에는 [4, 6, 5, 1]이 있습니다. 두 큐의 합은 30이므로 각 큐의 합은 15가 되어야 합니다. 일단 현재는 queue1의 합은 14, queue2의 합은 16입니다. 이때 queue1의 합이 queue2의 합보다 작음을 알 수 있죠. 이때 queue2의 첫번째 원소인 4를 queue1으로 가져옵니다. (queue1 합 : 18, queue2 합 : 12) queue1의 합이 queue2의 합보다 커졌습니다. 이제 반대로 queue1의 첫번째 원소인 3을 queue2로 가져옵니다. (queue1 합  = queue2 합 = 16)

 

여기서 규칙을 파악할 수 있습니다. queue1의 합과 queue2의 합을 비교하여, 더 큰 합을 가진 큐에서 다른 큐로 첫번째 원소를 넘겨준다는 것이죠. 

 

또한 두 큐의 합을 미리 계산하면 반복 횟수를 줄일 수 있습니다. 만약 두 큐의 합이 홀수라면, 어떤 경우에도 각 큐의 합이 서로 같아질 수 없습니다. 여기까지 반영하여 코드를 작성하면 아래와 같습니다.

 

 

from collections import deque
def solution(queue1, queue2):
    total = sum(queue1) + sum(queue2)
    if total % 2 : return -1
    
    answer = 0
    Q1,Q2 = deque(queue1),deque(queue2)
    S1,S2 = sum(Q1),sum(Q2)
    total_div = total // 2
    
    while S1 != S2 :
        S1,S2 = sum(Q1),sum(Q2)
            
        if S1 < S2 :
            Q1.append(Q2.popleft())
            if len(Q2) == 0 : 
                answer = -1
                break
            answer += 1
            
        elif S1 > S2:
            Q2.append(Q1.popleft())
            if len(Q1) == 0 : 
                answer = -1
                break
            answer += 1
            
        else :
            return answer

    return -1

 

 

그런데 위와 같이 코드를 작성하니, 테스트는 통과하나 제출 시 시간 초과가 다수 발생합다. 반복 횟수가 너무 많다는 것인데, 어디가 문제일까요?

 

먼저 발견한 점은 "가장 많이 반복하는 경우"를 생각할 수 있다는 것이었습니다. 두 큐의 합이 서로 계속 같지 않아 pop과 insert를 반복하면 두 큐의 원소가 서로 완전히 바뀌는 경우가 발생하고, 이후로 더 반복하면 다시 처음 상태로 돌아오게 됩니다. 이러한 경우는 모두 queue1(혹은 queue2)의 길이 * 3회에 수렴합니다. 이에 따라 while 반복문이 아닌 for 반복문으로 수정할 수 있습니다.

 

 

for _ in range(len(queue1)*3) :
    S1,S2 = sum(Q1),sum(Q2)

    if S1 < S2 :
        Q1.append(Q2.popleft())
        if len(Q2) == 0 : 
            answer = -1
            break
        answer += 1

    elif S1 > S2:
        Q2.append(Q1.popleft())
        if len(Q1) == 0 : 
            answer = -1
            break
        answer += 1

    else :
        return answer

 

 

그런데 웬걸, 여전히 시간 초과가 발생합니다. 여기서 무엇을 더 줄일 수 있을까요? 

 

해답은 바로 sum() 함수에 있었습니다. 위 코드를 보시면 반복이 한 번 일어날 때마다 Q1과 Q2의 sum을 구하고 있습니다. sum()은 O(n) 시간만큼 작업이 소요되기 때문에 Q1 혹은 Q2 안에 원소가 많다면 그만큼 sum()을 구하는 시간 또한 길어질 것입니다.

 

이에 따라 코드를 아래와 같이 수정할 수 있습니다.

 

 

from collections import deque
def solution(queue1, queue2):
    total = sum(queue1) + sum(queue2)
    if total % 2 : return -1
    
    answer = 0
    Q1,Q2 = deque(queue1),deque(queue2)
    S1,S2 = sum(Q1),sum(Q2)
    total_div = total // 2
    
    for _ in range(len(queue1)*3) :    
        if S1 < S2 :
            S1 += Q2[0]
            S2 -= Q2[0]
            Q1.append(Q2.popleft())
            answer += 1
            
        elif S1 > S2:
            S1 -= Q1[0]
            S2 += Q1[0]
            Q2.append(Q1.popleft())
            answer += 1
            
        else :
            return answer

    return -1