Fish – Codility – Solution

Java solution to Codility Fish problem (Lesson 7 – Stacks and Queues) which scored 100%. The problem is to determine how many fish are alive in a river of fish moving upstream and downstream.

The strategy is to keep a running stack of fish as go through array of fish swimming in same direction or upstream until encounter fish going in opposite direction and then figure out what fish will survive.


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
package com.codility.lesson07.stacksqueues;

import java.util.Stack;

public class Fish {
  public int solution(int[] A, int[] B) {
    final int UPSTREAM = 0;
   
    Stack fishStack = new Stack();
    fishStack.push(new Fish1(A[0], B[0]));

    for(int i=1; i<A.length; i++) { if(B[i] == fishStack.peek().direction) { //same direction, so put another fish on stack fishStack.push(new Fish1(A[i], B[i])); } //if top of stack fish is upstream, not right condition to see who's eating who else if(fishStack.peek().direction == UPSTREAM) { fishStack.push(new Fish1(A[i], B[i])); } else { //figure out who's eating who while(!fishStack.isEmpty()) { //current fish is swimming in same direction as top of stack fish - they can't eat other if(fishStack.peek().direction == B[i]) { fishStack.push(new Fish1(A[i], B[i])); break; } else { //existing fish that was on stack eats latest fish if(fishStack.peek().size > A[i]) {
              break; //eating finished
            }
            else {
              fishStack.pop();
              continue; //keep checking other elements on stack
            }
          }
        }
        if(fishStack.isEmpty()) {
          fishStack.push(new Fish1(A[i], B[i])); //current fish ate all the fish in the stack
        }
      }
    }
    return fishStack.size();
  }
 
  class Fish1 {
    private int size;
    private int direction;
   
    Fish1(int pSize, int pDirection) {
      size = pSize;
      direction = pDirection;
    }
  }
}

TestNG test cases for this problem which all passed:


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
package test.com.codility.lesson07.stacksqueues;

import org.testng.Assert;
import org.testng.annotations.BeforeTest;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

import com.codility.lesson07.stacksqueues.Fish;

public class FishTests {
  private Fish solution;
 
  @BeforeTest
  public void setUp() {
    solution = new Fish();
  }

  @DataProvider(name = "test1")
  public Object [][] createData1() {
    return new Object [][] {
      new Object [] {  new int [] { 4, 3, 2, 1, 5 }, new int [] { 1, 0, 1, 0, 1 },  3 },
      new Object [] {  new int [] { 4, 3, 2, 0, 5 }, new int [] { 0, 1, 0, 0, 0 },  2 },
      new Object [] {  new int [] { 4, 3, 2, 1, 5 }, new int [] { 0, 1, 0, 0, 0 },  2 },
      new Object [] {  new int [] { 4, 3, 2, 1, 5 }, new int [] { 0, 1, 1, 0, 0 },  2 },
      new Object [] {  new int [] { 4, 3, 2, 5, 6 }, new int [] { 1, 0, 1, 0, 1 },  2 },
      new Object [] {  new int [] { 7, 4, 3, 2, 5, 6 }, new int [] { 0, 1, 1, 1, 0, 1 },  3 },
      new Object [] {  new int [] { 3, 4, 2, 1, 5 }, new int [] { 1, 0, 0, 0, 0 },  4 },
      new Object [] {  new int [] { 3 }, new int [] { 1 },  1 },
      new Object [] {  new int [] { 3 }, new int [] { 0 },  1 },
    };
  }

  @Test(dataProvider = "test1")
  public void verifySolution(int [] pA, int [] pB, int pExpected) {  
    Assert.assertEquals(solution.solution(pA, pB), pExpected);
  }
}