stream 优化递归

需求

有一个部门类，如下

public class Department {
    /**
     * 部门 id
     */
    private String id;

    /**
     * 部门名称
     */
    private String name;

    /**
     * 上级部门id， 顶级为 null
     */
    private String parentId;

    ... 省略 get set
}

有一组如下的数据需要进行排序，排序规则为，1级部门后面是所有2级部门，如果2级部门存在下属3级部门，则该2级部门后紧跟着所有3级部门，以此类推。

List<Department> originalList = new ArrayList<>();
originalList.add(new Department("7", "国内采购部", "5"));
originalList.add(new Department("8", "北上广采购部", "7"));
originalList.add(new Department("9", "沿海采购部", "7"));
originalList.add(new Department("2", "国外采购部", "5"));
originalList.add(new Department("10", "欧美采购部", "2"));
originalList.add(new Department("11", "澳洲采购部", "2"));
originalList.add(new Department("3", "国内销售部", "6"));
originalList.add(new Department("4", "国外销售部", "6"));
originalList.add(new Department("1", "运营部", null));
originalList.add(new Department("5", "采购部", "1"));
originalList.add(new Department("6", "销售部", "1"));

排序结果如下

排序结果.png

简单版

实现分析

• 找到所有一级部门
• 遍历所有一级部门，找到是否存在二级部门，不存在返回
• 遍历所有二级部门，查找是否存在三级部门，不存在返回
• 遍历所有三级部门，查找是否存在四级部门，不存在返回
• ...
  能看出来这是一个递归算法，尝试用伪代码实现下逻辑

List<Department> 排序结果;

void 查找子部门(上级部门) {
  所有下级部门 = 查找所有下级部门()
  if (所有下级部门不存在) {
    return;
  }
  for (下级部门: 所有下级部门) {
    排序结果.添加(下级部门);
    查找子部门(下级部门);
  }
}

思路清楚了，就是找下级，遍历下级，添加下级到结果集中，同时找下级的所有下级

代码实现

public class DepartmentSort {

    private List<Department> deptList;

    private List<Department> resultList = new ArrayList<>();

    public DepartmentSort(List<Department> deptList) {
        this.deptList = deptList;
    }

    public static List<Department> sort(List<Department> originalList) {
        return new DepartmentSort(originalList).sort();
    }

    private List<Department> sort() {
        findChildren(new Department());
        return resultList;
    }

    /**
     * 查询下级部门
     *
     * @param parentDept
     */
    private void findChildren(Department parentDept) {
        // 查找所有下级部门
        List<Department> childrenList = deptList.stream()
                .filter(d -> Objects.equals(parentDept.getId(), d.getParentId()))
                .collect(Collectors.toList());

        if (childrenList.isEmpty())
            /* 不存在下级部门，跳出递归 */
            return;
        else
            // 遍历所有下级部门，塞入下级部门 resultList 的同时，查找所有该下级部门对应的下级部门
            childrenList.forEach(d -> {
                resultList.add(d);
                findChildren(d);
            });
    }

}

编码测试

List<Department> resultList = DepartmentSort.sort(originalList);

测试结果，达到了预期结果

测试结果1.png

优化版

接下来，分析上个版本代码的不足

• 上个版本的实现思路是用 一个 resultList 集合去存储所有结果， findChildren 方法是返回 void 的
• 但其实对于上级部门，调用 findChildren 方法，希望得到的是所有下级部门以及下级部门可能存在的所有更下一级部门的集合

带着这几条思路，改造代码如下

private List<Department> findChildren2(Department parentDept) {
    // 查找所有下级部门
    List<Department> childrenList = deptList.stream()
            .filter(d -> Objects.equals(parentDept.getId(), d.getParentId()))
            .collect(Collectors.toList());

    List<Department> tempList = new ArrayList<>();
    // 遍历所有下级
    childrenList.forEach(d -> {
        tempList.add(d);
        tempList.addAll(findChildren2(d));
    });
    return tempList;
}

经过测试，结果正确，不贴图了...

stream 优化版

继续分析上个版本代码，依然存在不优雅的地方

• 需要弄一个 tempList 存储所有下级部门以及可能存在的所有更下一级部门
• 虽然使用了 stream，但是只是用来过滤下级，大材小用了

利用 stream 流的思想去考虑问题，其实遍历出来的下级和可能存在的所有下级部门的所有下级部门，是绑定在一起的，也就是 tempList。根据这样的关系，很容易想到 map 操作符，得到 下级部门.map(下级部门和所有下级部门的下级部门的集合)

带着这个思路，实现如下代码

public class DepartmentSort2 {

    private List<Department> deptList;

    public DepartmentSort2(List<Department> deptList) {
        this.deptList = deptList;
    }

    public static List<Department> sort(List<Department> originalList) {
        return new DepartmentSort2(originalList).sort();
    }

    private List<Department> sort() {
        return findChildren(new Department())
                .collect(Collectors.toList());
    }

    /**
     * 查询下级部门
     *
     * @param parentDept
     */
    private Stream<Department> findChildren(Department parentDept) {
        return deptList.stream()
                .filter(d -> Objects.equals(parentDept.getId(), d.getParentId()))
                .flatMap(d -> Stream.concat(Stream.of(d), findChildren(d)));
    }

}

编写测试代码

List<Department> resultList2 = DepartmentSort2.sort(originalList);

测试结果

stream 版测试结果.png

源码 https://github.com/MoonBottle/sortdemo