Identifying arrays with the same mean

Question

I came across a problem where given a array of integer arrays of different lengths [[1,2,3],[4,1,1], [9,2,1]] you need to return an array of arrays, each array containing indices of the arrays (of the input array) such that the corresponding arrays have the same mean: [[0,1],[2]] This seems relatively simple to solve using Python:

def groupByMean(a):
    d,e=[],[]
    for i,j in enumerate(a):
        if sum(j)/len(j)not in e:
            e+=[sum(j)/len(j)]
            d+=[[i]]
        else:
            d[e.index(sum(j)/len(j))]+=[i]
    return d

However, when trying to solve this in Java this was my approach: using a hashmap, map each new mean to a list of the corresponding indices. Then iterate the hashmap, to get the arraylists and convert them to int[] arrays and construct the 2d array ...

Is there a simpler approach to solve this using Java?

This is my java code - looking for a different way to solve this:

public static void main(String[] args) {
    int[][] arr = { { 1, 2, 3 }, { 2, 3, 4 }, { 2, 4, 0 } };
    for (int[] nums : groupBySum(arr)) {
        for (int n : nums) {
            System.out.print(n + " ");
        }
        System.out.println();
    }
}

public static int[][] groupByMean(int[][] arr) {
    Map<Double, List<Integer>> map = new HashMap<>();
    int i = 0;
    for (int[] nums : arr) {
        double average = getAverage(nums);
        if (!map.containsKey(average)) {
            List<Integer> indices = new ArrayList<>();
            indices.add(i);
            map.put(average, indices);
        } else {
            map.get(average).add(i);
        }
        i++;
    }
    int[][] result = new int[map.size()][];
    int row = 0;
    for (List<Integer> indices : map.values()) {
        result[row] = new int[indices.size()];
        for (int k = 0; k < indices.size(); k++) {
            result[row][k] = indices.get(k);
        }
        row++;
    }
    return result;
}

public static double getAverage(int[] arr) {
    int sum = 0;
    for (int num : arr) {
        sum += num;
    }
    return ((double) sum) / arr.length;
}

Answer 1

Nice implementation. Few suggestions to make the method groupByMean more compact using Java Streams:

• Calculate the average:

  public static double getAverage(int[] arr) {
    int sum = 0;
    for (int num : arr) {
        sum += num;
    }
    return ((double) sum) / arr.length;
  }
  

  To:

  public static double getAverage(int[] arr) {
    return Arrays.stream(nums).average().getAsDouble();
  }
  

• Group by average:

  if (!map.containsKey(average)) {
      List<Integer> indices = new ArrayList<>();
      indices.add(i);
      map.put(average, indices);
  } else {
      map.get(average).add(i);
  }
  

  To:

  map.computeIfAbsent(average, v -> new ArrayList<>()).add(i);
  

• Convert list to array:

  for (int k = 0; k < indices.size(); k++) {
      result[row][k] = indices.get(k);
  }
  

  To:

  result[row] = indices.stream().mapToInt(index->index).toArray();
  

• Convert map's values to a matrix:

  int[][] result = new int[map.size()][];
  int row = 0;
  for (List<Integer> indices : map.values()) {
      result[row] = new int[indices.size()];
      for (int k = 0; k < indices.size(); k++) {
          result[row][k] = indices.get(k);
      }
      row++;
  }
  return result;
  

  To:

  return map.values().stream()
            .map(v -> v.stream().mapToInt(index->index).toArray())
            .toArray(int[][]::new);
  

Time/Space Complexity

In terms of complexity there are no relevant differences between your solution and the one using Streams. The time complexity is still \$O(n*m)\$ where \$n\$ is the number of arrays and \$m\$ is the size of the longest array. Basically, for each array we need to calculate the average.

To check which approach is faster you need to benchmark the solutions.

Final code

public static void main(String[] args) {
    int[][] arr = {{ 1, 2, 3 }, { 2, 3, 4 }, { 2, 4, 0 }};
    Arrays.stream(groupByMean(arr)).map(Arrays::toString)
        .forEach(System.out::println);
}

public static int[][] groupByMean(int[][] arr) {
    Map<Double, List<Integer>> map = new HashMap<>();
    for (int i=0 ; i<arr.length; i++) {
        double average = Arrays.stream(arr[i]).average().getAsDouble();
        map.computeIfAbsent(average, v -> new ArrayList<>()).add(i);
    }
    return map.values().stream()
      .map(v -> v.stream().mapToInt(index->index).toArray())
      .toArray(int[][]::new);
}

Answer 2

I would suggest using the groupingBy collector.

public static void main(String[] args) {
    int[][] arr = {{ 1, 2, 3 }, { 2, 3, 4 }, { 2, 4, 0 }};
    IntStream.range(0, arr.length)
        .boxed()
        .collect(groupingBy(i->Arrays.stream(arr[i]).average().getAsDouble()))
        .values()
        .forEach(System.out::println);
}