[Codility] Stacks and Queues - Fish 풀이

예문


https://app.codility.com/programmers/lessons/7-stacks_and_queues/fish/

You are given two non-empty arrays A and B consisting of N integers. Arrays A and B represent N voracious fish in a river, ordered downstream along the flow of the river.

The fish are numbered from 0 to N − 1. If P and Q are two fish and P < Q, then fish P is initially upstream of fish Q. Initially, each fish has a unique position.

Fish number P is represented by A[P] and B[P]. Array A contains the sizes of the fish. All its elements are unique. Array B contains the directions of the fish. It contains only 0s and/or 1s, where:

0 represents a fish flowing upstream, 1 represents a fish flowing downstream. If two fish move in opposite directions and there are no other (living) fish between them, they will eventually meet each other. Then only one fish can stay alive − the larger fish eats the smaller one. More precisely, we say that two fish P and Q meet each other when P < Q, B[P] = 1 and B[Q] = 0, and there are no living fish between them. After they meet:

If A[P] > A[Q] then P eats Q, and P will still be flowing downstream, If A[Q] > A[P] then Q eats P, and Q will still be flowing upstream. We assume that all the fish are flowing at the same speed. That is, fish moving in the same direction never meet. The goal is to calculate the number of fish that will stay alive.

For example, consider arrays A and B such that:

  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

Initially all the fish are alive and all except fish number 1 are moving upstream. Fish number 1 meets fish number 2 and eats it, then it meets fish number 3 and eats it too. Finally, it meets fish number 4 and is eaten by it. The remaining two fish, number 0 and 4, never meet and therefore stay alive.

Write a function:

class Solution { public int solution(int[] A, int[] B); }

that, given two non-empty arrays A and B consisting of N integers, returns the number of fish that will stay alive.

For example, given the arrays shown above, the function should return 2, as explained above.

Write an efficient algorithm for the following assumptions:

N is an integer within the range [1..100,000];
each element of array A is an integer within the range [0..1,000,000,000];
each element of array B is an integer that can have one of the following values: 0, 1;
the elements of A are all distinct.

해석

  • A, B : 비어있지 않은 N개의 정수로 구성 된 배열
  • 물고기 : 0 ~ N - 1
  • 물고기 P, Q가 있으면 : P < Q
  • P : 물고기 번호. A[P], B[P]
  • 배열 A : 물고기 size
  • 배열 B : 물고기의 방향
    • 0 : 상류
    • 1 : 하류
  • 두 물고기가 서로 다른 방향으로 움직이고 서로 만날 때 큰 물고기가 작은 물고기 잡아먹음
  • P < Q, B[P] = 1 이고 B[Q] = 0 일 때 두 물고기는 서로 만난 후
    • A[P] > A[Q] 이면 P가 Q 잡아먹고 P는 하류로 이동함.
    • A[Q] > A[P] 이면 Q가 P 잡아먹고 Q는 상류로 이동함.
  • 같은 방향으로 움직이는 물고기는 서로 만나지 않음.
  • 살아있는 물고기의 수를 반환하라!

풀이

  • 하류로 이동하는 물고기만 Stack에 담는다.
  • 상류로 이동하는 물고기를 만날 때 사이즈 비교
    • Stack이 비어 있지 않으면 사이즈 비교를 하고 상류로 이동하는 물고기가 더 클 경우 Stack에서 pop을 해준다.
    • 상류로 올라가는 경우가 하류로 내려오는 물고기들 보다 클 경우 Stack은 비게 되고 상류로 올라가는 물고기 수가 증가된다.

제약사항

  • N의 범위는 정수 [0..100,000]
  • 배열 A의 각 요소의 범위는 정수 [0..1,000,000,000]
  • 배열 B의 각 요소는 0, 1 중에 하나
  • A의 모든 요소는 유니크

코드

import java.util.Stack;
public int solution(int[] A, int[] B) {
  Stack<Integer> down = new Stack<>();
  int lastSize;
  int aliveCount = 0;

  for (int i = 0; i < A.length; i++) {
    if (B[i] == 1) down.push(A[i]);
    else {
      while (!down.isEmpty()) {
        lastSize = down.peek();
        if (lastSize > A[i]) break;
        else down.pop();
      }
      if (down.isEmpty()) aliveCount++;
    }
  }
  return aliveCount + down.size();
}

결과