[알고리즘] Fishes 살아남은 물고기 세기 문제

내 풀이는 다음과 같다. 맞긴 했는데 조금 까다로웠던 것 같아서 더 좋은 풀이가 있는지 궁금했다.

# you can write to stdout for debugging purposes, e.g.
# print("this is a debug message")

# upstream: 0 (앞으로)
# downstream: 1 (뒤로)


#   A[0] = 4    B[0] = 0
#   A[1] = 3    B[1] = 1
#   A[2] = 2    B[2] = 0
#   A[3] = 1    B[3] = 0
#   A[4] = 5    B[4] = 0

# 먹히거나, 먹거나 둘 중 하나만 있음. 무승부 없음.

def solution(A, B):
    fishes:list[tuple(int, int)] = []
    for i in range(len(A)):
        fishes.append((A[i], B[i]))
    # print(fishes)

    # 만나는 물고기들만 승부한다.
    # 앞에서부터 하나씩 스택에 넣어서 확정하자.
    # 두 마리씩 비교해야 한다. 바로 이전 물고기는 stack의 top으로 찾고,
    # 새로운 물고기는 만약 방향이 upstream이면 그 이전 물고기와 비교해서 방향이 같은 물고기가 나올 때까지 승부한다. 방향이 같으면 stack에 넣고 비교 끝. 
    # 방향이 다르면 승부해서 한마리만 남김.
    stack:list[tuple(int, int)] = []
    for fish in fishes:
        # 일단 새로운 물고기를 스택에 추가
        stack.append(fish)
        # 맨 마지막에 있는 물고기와 그 이전 물고기들을 순서대로 승부 여부 판단.
        # 두 마리를 넣어서 승부하는 메소드 -> 한 마리만 남김
        # 남은게 한 마리가 아니고 마지막 두 마리가 다른 방향이면 계속 승부
        while len(stack)!=1 and stack[-2][1] != stack[-1][1] and stack[-1][1] !=1:
            # 승부
            if stack[-2][0] < stack[-1][0]:
                stack.pop(-2)
            else:
                stack.pop()
            # print(stack)
                
    # print(stack)
    # print(len(stack))
    return len(stack)

그래서 다음과 같이 다듬을 수 있다.

---

내 코드 스타일 유지 + 다듬은 버전

def solution(A, B):
    stack: list[tuple[int, int]] = []  # (size, dir)

    for size, dir_ in zip(A, B):
        stack.append((size, dir_))

        # 뒤가 0(up), 앞이 1(down)일 때만 계속 싸움
        while len(stack) >= 2 and stack[-2][1] == 1 and stack[-1][1] == 0:
            # 두 마리 중 큰 놈만 살아남음
            if stack[-2][0] < stack[-1][0]:
                # 앞(down)이 잡아먹힘
                stack.pop(-2)
            else:
                # 뒤(up)이 잡아먹힘
                stack.pop()

    return len(stack)

• list[tuple[int,int]] 문을 순회할 때에는 for a, b in zip(A,B) 를 쓰면 메모리 낭비도 줄어서 더 좋다고 한다.
• 승부 조건을 좀 더 명확하게 할 수 있었다. 두 물고기의 방향이 다르고 앞 물고기가 downstream 으로 헤엄친다라는 조건은, 그냥 앞 물고기는 downstream, 뒤 물고기는 upstream 이어야 한다로 바꾸면 더 명확하고 간결하다.

# you can write to stdout for debugging purposes, e.g.
# print("this is a debug message")

#   A[0] = 4    B[0] = 0
#   A[1] = 3    B[1] = 1
#   A[2] = 2    B[2] = 0
#   A[3] = 1    B[3] = 0
#   A[4] = 5    B[4] = 0

# 0 represents a fish flowing upstream,(앞으로)
# 1 represents a fish flowing downstream.(뒤로)
# 강을 올라가는 물고기, 내려가는 물고기가 있음.
# 두 물고기가 방향이 다르면, 크기 비교를 해서 한놈이 한놈을 잡아먹음.
# 스택으로 푸려고 함.
# 최종 살아남는 물고기를 upstream + downstream 으로 나눌 수 있음.
# downstream 물고기를 스택에 쌓고,
# 앞으로 가는(upstream) 물고기를 스택의 끝부터 비교해서 전체를 거쳐 살아남는지 보면 된다.


def solution(A, B):
    downstreamStack:list[int] = []
    alive:int = 0
    for size, direction in zip(A, B):
        if direction == 1: # downstream
            downstreamStack.append(size)
        else: # upstream
            # downstreamStack 모두와 맨 뒤부터 승부해서 살아남는지 보자.
            while downstreamStack:
                # 승부
                if downstreamStack[-1] < size:
                    # upstream Fish가 먹음
                    downstreamStack.pop()
                else:
                    break
            # downstreamStack이 비면 살아남은 것.
            if not downstreamStack:
                alive +&#x3D; 1
    return alive + len(downstreamStack)