求二维数组均值并按均值分组返回索引二维数组的实现问题

这是来自CodeSignal的问题：

问题描述 #

给定二维数组a，需将子数组a[i]按均值分组：

• 相同均值的子数组归为一组，不同均值分属不同组
• 每组包含对应子数组的索引（如i、j等）
• 返回的分组为二维数组，要求：
  1. 每组内索引升序排列
  2. 各组按组内最小索引升序排列

输入示例 #

int[][] a = {{3, 3, 4, 2}, {4, 4}, {4, 0, 3, 3}, {2, 3}, {3, 3, 3}};

输出示例 #

solution(a) = {{0, 4}, {1}, {2, 3}};

现有代码 #

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; 
}

遇到的问题 #

不知道如何将均值数组中的索引和对应均值存入ArrayList，尝试过for循环和HashMap但未成功。

解决思路与修正后的代码 #

首先修正均值计算的错误 #

现有代码中，内部循环（j的循环）里每次累加后都计算均值并赋值，会导致均值被多次覆盖，最终得到未累加完的错误结果。正确做法是先累加完当前子数组的所有元素，再计算均值。

使用HashMap分组 #

用HashMap<Integer, List<Integer>>存储均值与对应索引列表的映射：

• Key：子数组的均值
• Value：该均值对应的所有子数组索引组成的列表

排序规则处理 #

1. 按索引顺序遍历添加的列表，内部索引天然升序，若需稳妥可显式排序。
2. 最后将所有分组按组内最小索引升序排列，排序依据为每组的第一个元素（即最小索引）。

修正后的完整代码 #

import java.util.*;
import java.util.stream.Collectors;

public static int[][] sortMean(int[][] a) {
    // 用HashMap存储均值对应的索引列表
    Map<Integer, List<Integer>> meanToIndices = new HashMap<>();
    
    for (int i = 0; i < a.length; i++) {
        int sum = 0;
        // 累加当前子数组的所有元素
        for (int num : a[i]) {
            sum += num;
        }
        // 计算均值（题目中为整数除法，与示例逻辑一致）
        int mean = sum / a[i].length;
        
        // 将当前索引加入对应均值的列表，不存在则新建列表
        meanToIndices.computeIfAbsent(mean, k -> new ArrayList<>()).add(i);
    }
    
    // 将分组转换为列表并按规则排序
    List<List<Integer>> groups = meanToIndices.values().stream()
            // 显式确保每组内索引升序
            .map(list -> {
                Collections.sort(list);
                return list;
            })
            // 按组内最小索引升序排列各组
            .sorted(Comparator.comparingInt(list -> list.get(0)))
            .collect(Collectors.toList());
    
    // 转换为二维int数组返回
    return groups.stream()
            .map(list -> list.stream().mapToInt(Integer::intValue).toArray())
            .toArray(int[][]::new);
}

代码说明 #

1. computeIfAbsent方法：自动处理均值不存在的情况，无需手动判断，直接添加索引到对应列表。
2. 分组排序：通过sorted(Comparator.comparingInt(list -> list.get(0)))实现按每组最小索引排序。
3. 流操作转换：最后将List<List<Integer>>转换为int[][]类型，完全符合题目返回要求。

内容的提问来源于stack exchange，提问作者Yunmi Lee