[자료구조] [python] Equal Stacks

Hackerrank의 Data structures 관한 문제 중에서 Easy에 속하는 문제이다. 원 문제 링크는 다음을 따라가보자. 

www.hackerrank.com/challenges/equal-stacks/problem

Equal Stacks | HackerRank

 

 

• 문제의 하단에 문제에 서술된 상황을 묘사하는 그림이 있어서, 이해하는 데 어렵지 않았다. 문제를 풀이하는 데 있어서 'stack' 자료구조를 활용하는 것이 도움이 되었다. 스택은 배열의 끝에서만 데이터를 접근할 수 있는 선형 자료구조다. 파이썬에서는 append()와 pop() 메서드로 스택을 구현한다. 

 

• 아래 그림의 h1 실린더의 블록은 [3,2,1,1,1]로 구성되어 있다. 가장 위의 블록이 맨 위에 있는 구조다. h1,h2,h3의 실린더의 블록들을 차례로 훑으며, stack1, stack2, stack3에 (실린더의 가장 위의 블록의 크기, 해당 블록을 제거했을 때의 남은 실린더의 높이)의 튜플형태로 stack구조로 만들었다. 또한, 아래서부터 블록을 훑으면서 블록의 높이를 계산해야 중복계산하지 않게 되므로, 인자로 받은 h1 리스트를 뒤 원소부터 빼면서 높이를 계산했다. 

• 유의해야 할 점 : while문 안에  max_value = min([sum1,sum2,sum3])로 하면, 2개의 테스트 케이스에서 runtimeerror가 나온다. 이에, 최소값을 설정해놓고 특정 실린더에서 블록을 하나씩 제거할 때마다 if문으로 최소값인지 확인하며, 업데이트해주었다. 특히, 이 부분은 다른 문제에서도 빈출로 다뤄진다.

import math
import os
import random
import re
import sys

#
# Complete the 'equalStacks' function below.
#
# The function is expected to return an INTEGER.
# The function accepts following parameters:
#  1. INTEGER_ARRAY h1
#  2. INTEGER_ARRAY h2
#  3. INTEGER_ARRAY h3
#


def equalStacks(h1, h2, h3):
    stack1, stack2, stack3 = [],[],[]
    sum1, sum2, sum3 = 0,0,0

    
    #(블록, 블록을 빼고 나서의 height)
    
    #리스트의 원소가 있을 경우에만 while 문 안에 있는 것을 처리한다. 
    
    while h1:
    	#리스트의 마지막 원소 일 수록, 실린더의 하단에 위치한 블록이므로 뒤 원소부터 빼면서 전체 높이를 구한다.
        h = h1.pop()
        stack1.append((h, sum1))
        sum1 += h
        
    while h2:
        h = h2.pop()
        stack2.append((h, sum2))
        sum2 += h
    
        
    while h3:
        h = h3.pop()
        stack3.append((h, sum3))
        sum3 += h
    
    max_value = min([sum1,sum2,sum3])
    
    #세 실린더의 높이가 동일할 때까지 while문 내의 코드를 동작한다.
    
    while not (sum1 == sum2 == sum3):
   
        
        if max_value < sum1:
            _, sum1 = stack1.pop()
            
            #세 높이를 맞춰줘야 하므로, 가장 높은 블록을 제거한 높이가 최소 높이인지 체크해야 한다. 
            #while 문 max_value = min([sum1,sum2,sum3])로 하면, 테스트 케이스의 몇 개의 문제가 
            #runtime error가 나온다. 
            if sum1 < max_value:
                max_value = sum1
        
        if max_value < sum2:
            _, sum2 = stack2.pop()
            if sum2 < max_value:
                max_value = sum2
        
        if max_value < sum3:
            _, sum3 = stack3.pop()
            if sum3 < max_value:
                max_value = sum3
            
    
    return max_value
    # Write your code here

if __name__ == '__main__':
    fptr = open(os.environ['OUTPUT_PATH'], 'w')

    first_multiple_input = input().rstrip().split()

    n1 = int(first_multiple_input[0])

    n2 = int(first_multiple_input[1])

    n3 = int(first_multiple_input[2])

    h1 = list(map(int, input().rstrip().split()))

    h2 = list(map(int, input().rstrip().split()))

    h3 = list(map(int, input().rstrip().split()))

    result = equalStacks(h1, h2, h3)

    fptr.write(str(result) + '\n')

    fptr.close()