The challenge

Write a method, that will get an integer array as parameter and will process every number from this array.
Return a new array with processing every number of the input-array like this:

1
2
3

If the number has an integer square root, take this, otherwise square the number.

[4,3,9,7,2,1] -> [2,9,3,49,4,1]

The input array will always contain only positive numbers and will never be empty or null.

Test cases

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

import org.junit.Test;
import static org.junit.Assert.assertEquals;
import org.junit.runners.JUnit4;
import java.util.Arrays;

public class SolutionTest {
    @Test
    public void basicTests() {
        int[] input = new int[] { 4, 3, 9, 7, 2, 1 };
        int[] expected = new int[] { 2, 9, 3, 49, 4, 1 };
        assertEquals(Arrays.toString(expected), Arrays.toString(Squarer.squareOrSquareRoot(input)));

        input = new int[] { 100, 101, 5, 5, 1, 1 };
        expected = new int[] { 10, 10201, 25, 25, 1, 1 };
        assertEquals(Arrays.toString(expected), Arrays.toString(Squarer.squareOrSquareRoot(input)));

        input = new int[] { 1, 2, 3, 4, 5, 6 };
        expected = new int[] { 1, 4, 9, 2, 25, 36 };
        assertEquals(Arrays.toString(expected), Arrays.toString(Squarer.squareOrSquareRoot(input)));        
    }
}

The solution in Java

Option 1 (using a loop):

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

public class Squarer {
  public static int[] squareOrSquareRoot(int[] array) {
    
    for (int i=0; i<array.length; i++) {
        
        double sqr = Math.sqrt(array[i]);
        if (sqr==Math.round(sqr)) {
            // can square
            array[i] = (int)sqr;
        } else {
            // perform square
          array[i] = array[i]*array[i];
        }
        
    }
    
    return array;
  }
}

Option 2 (using a stream):

1
2
3
4
5
6
7
8
9

import java.util.Arrays;

public class Squarer {
  public static int[] squareOrSquareRoot(int[] array) {
    return Arrays.stream(array)
          .map(i -> Math.sqrt(i) % 1 == 0 ? ((int) Math.sqrt(i)) : (i * i))          
          .toArray();      
  }
}

Option 3 (using an IntStream):

1
2
3
4
5
6
7

import static java.util.stream.IntStream.of;

class Squarer {
  static int[] squareOrSquareRoot(int[] array) {
    return of(array).map(i -> Math.sqrt(i) % 1 == 0 ? (int) Math.sqrt(i) : i * i).toArray();
  }
}