출처
https://app.codility.com/programmers/lessons/7-stacks_and_queues/fish/
개인적으로 헷갈렸던 문제여서 포스팅을 작성해봅니다.
문제
강 한줄기가 존재한다고 할 때, 위와 아래로 이동하는 물고기들의 방향과 크기의 집합이 주어진다. 이때 서로 마주치게 되는 물고기는 크기가 큰 물고기가 작은 물고기를 잡아먹게 되며, 최종적으로 살아남은 물고기의 수를 구하는 문제이다.
풀이
처음에는 강 한줄기라고 생각해 파이썬의 deque를 이용해 접근하였지만, 효율성이 매우 떨어졌고 해결하지 못해 Stack을 이용한 풀이로 바꾸었다.
강 한줄기 안에서 통행하기 때문에, 물고기가 만나는 경우는 오직 방향이 반대인 경우(0, 1)밖에 존재하지 않는다. 방향이 들어있는 B 배열을 순회하며 방향이 1인 경우는 스택에 크기를 넣어주고, 그렇지 않은 경우 while 반복문을 통해 크기를 비교해준다.
만약 스택에 맨 위에 있는 방향이 1인 물고기의 크기가 현재 물고기의 크기보다 크다면, 현재 물고기는 잡아먹히게 되고 pop()했던 방향이 1인 물고기를 다시 append()를 통해 넣어주고 while문을 종료한다.
스택 맨 위 물고기 보다 현재 물고기가 더 크다면, 방향이 1인 물고기는 잡아먹히게 되며, 스택이 빌 때 까지 반복한다. 방향이 1인 물고기가 스택에 들어있지 않다면, 이 물고기는 무조건 살아남는다고 할 수 있기 때문에 살아남은 물고기수에 +1을 해준다.
마지막으로 스택에 남아있는 물고기가 있다면 그 수만큼 살아남은 물고기수에 더해주면 정답을 구할 수 있다.
코드
def solution(A, B):
stack = []
answer = 0
for i in range(len(A)):
if B[i] == 1:
stack.append(A[i])
else:
while stack:
tmp = stack.pop()
if tmp > A[i]:
stack.append(tmp)
break
if not stack:
answer += 1
answer += len(stack)
return answer