如何使用Java Streams提取嵌套列表各子列表第n个元素并收集为数组

问题描述 #

现有如下函数，接收嵌套ArrayList参数allLists和索引参数nth，返回所有子列表的第n个元素：
示例：当allLists = {{1,2,3}, {4,5,6}, {7,8,9}}、nth = 1时，函数返回{2,5,8}
原实现代码：

public static String[] getEveryNthElement(ArrayList<ArrayList<String>> allLists, int nth) {
        String[] nthList = new String[allLists.size()];
        int n = 0;
        for (ArrayList<String> sList: allLists) {
            if (nth <= sList.size()) {
                nthList[n] = (sList.get(nth));
            }
            n += 1;
        }
        return nthList;
    }

已经通过Java Streams实现了可打印输出对应结果的版本，代码如下：

group.stream()
    .forEach(items -> {
        System.out.println(items.get(1));  // prints 2, 5, 8
});

问：如何将上述Streams处理得到的结果直接收集为数组类型？

解答 #

你只需要将forEach遍历打印的逻辑替换为map映射+toArray收集即可，代码如下：

// 兼容Java 11及以上版本
String[] nthList = group.stream()
        // 保留原逻辑的同时修正了边界判断，避免索引越界
        .map(sList -> nth < sList.size() ? sList.get(nth) : null)
        .toArray(String[]::new);

如果使用的是Java 8版本，调整最后一行的数组生成写法即可：

// 兼容Java 8及以上版本
String[] nthList = group.stream()
        .map(sList -> nth < sList.size() ? sList.get(nth) : null)
        .toArray(size -> new String[size]);

注意：你原本的for循环中边界判断存在问题，nth <= sList.size()会在nth等于子列表长度时触发IndexOutOfBoundsException，因为List的索引范围是0 ~ size-1，上述Stream实现已经修正了这个问题，你也可以同步调整原for循环的判断条件。

上述实现和你预期的逻辑完全一致，示例中传入nth=1时，得到的结果就是["2","5","8"]。

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