Get mean value from 2D arrays and return the 2D array with indices of same values-CodePudding

This is the problem from CodeSignal (http://codesignal.com).

Question:

You are given an array of arrays a. Your task is to group the arrays a[i] by their mean values, so that arrays with equal mean values are in the same group, and arrays with different mean values are in different groups.

Each group should contain a set of indices (i, j, etc), such that the corresponding arrays (a[i], a[j], etc) all have the same mean. Return the set of groups as an array of arrays, where the indices within each group are sorted in ascending order, and the groups are sorted in ascending order of their minimum element.

So if the input is:
 a = [[3, 3, 4, 2], [4, 4], [4, 0, 3, 3], [2, 3], [3, 3, 3]]
 The output should be:
solution(a) = [[0, 4], [1], [2, 3]]

Following is the code I did so far,

public static int[][] sortMean(int[][] a) {
    int[] means = new int[a.length]; //store means example: [3, 4, 2, 2, 3] 
    
    //get mean value from each array
    //store the values in mean array 
    for (int i=0; i<a.length; i  ) {
        int sum = 0; 
        for (int j=0; j<a[i].length; j  ) {
            sum  = a[i][j]; 
            int mean = sum / a[i].length;
            means[i] = mean; 
        }
    }
    
    ArrayList<List<Integer>> groupList = new ArrayList<>(); 
    
    //arraylist to 2d-array
    int[][] b = groupList.stream().map(  u  ->  u.stream().mapToInt(i->i).toArray()  ).toArray(int[][]::new);
    return b; 
}

The problem I'm facing is that I don't know how to store the indices and values from the mean array to ArrayList.

I tried for loops and HashMap but it did not work.

CodePudding user response：

You can use the streams API to group each array by the average value and convert it to a two dimensional array.

static int[][] sortMean(int[][] a) {
    return IntStream.range(0, a.length)
        .boxed()
        .collect(Collectors.groupingBy(
            i -> IntStream.of(a[i]).average().getAsDouble(),
            LinkedHashMap::new, Collectors.toList()))
        .values().stream()
        .map(list -> list.stream().mapToInt(Integer::intValue).toArray())
        .toArray(int[][]::new);
}

public static void main(String[] args) {
    int[][] a =  {{3, 3, 4, 2}, {4, 4}, {4, 0, 3, 3}, {2, 3}, {3, 3, 3}};
    int[][] output = sortMean(a);
    System.out.println(Arrays.deepToString(output));
}

output:

[[0, 4], [1], [2, 3]]