采用递归的方式转换数组为树形结构

项目快上线了，工作进入了尾声，下午在划水，花了半天的时间写了这个列表转树形结构的功能，采用递归的方式。以后需要用到的同学可以拿来跑跑，用到了阿里的fastjson。嗯，就这样

@Test
public void listToTreeTest() {
    // 有顶级节点
    String arrayJson1 = "[ {\n" +
            "\t\"id\": \"0001\",\n" +
            "\t\"pid\": \"0001\",\n" +
            "\t\"name\": \"总行事业部\"\n" +
            "}, {\n" +
            "\t\"id\": \"0002\",\n" +
            "\t\"pid\": \"0001\",\n" +
            "\t\"name\": \"上海分行\"\n" +
            "}, {\n" +
            "\t\"id\": \"0003\",\n" +
            "\t\"pid\": \"0001\",\n" +
            "\t\"name\": \"苏州分行\"\n" +
            "}, {\n" +
            "\t\"id\": \"0004\",\n" +
            "\t\"pid\": \"0002\",\n" +
            "\t\"name\": \"上海浦东支行\"\n" +
            "}, {\n" +
            "\t\"id\": \"0005\",\n" +
            "\t\"pid\": \"0003\",\n" +
            "\t\"name\": \"苏州吴江支行\"}]";

    // 无顶级节点
    String arrayJson2 = "[{\n" +
            "\t\"id\": \"0002\",\n" +
            "\t\"pid\": \"0001\",\n" +
            "\t\"name\": \"上海分行\"\n" +
            "}, {\n" +
            "\t\"id\": \"0003\",\n" +
            "\t\"pid\": \"0001\",\n" +
            "\t\"name\": \"苏州分行\"\n" +
            "}, {\n" +
            "\t\"id\": \"0004\",\n" +
            "\t\"pid\": \"0002\",\n" +
            "\t\"name\": \"上海浦东支行\"\n" +
            "}, {\n" +
            "\t\"id\": \"0005\",\n" +
            "\t\"pid\": \"0003\",\n" +
            "\t\"name\": \"苏州吴江支行\"}]";

    // 切换arrayJson1、arrayJson2测试
    JSONArray array = JSONArray.parseArray(arrayJson1);
    JSONArray result = listToTree(array, "id", "pid", "child", "0001");
    System.out.println(JSON.toJSONString(result));
}

/**
 * list转树结构list
 * @param array
 * @param idKey
 * @param pidKey
 * @param childKey
 * @param topName
 * @return
 */
public static JSONArray listToTree(JSONArray array, String idKey, String pidKey, String childKey, String topName) {
    if(hasTopNode(array, idKey, pidKey, topName)) {
        // 移除顶级元素
        JSONArray tempArray = new JSONArray();
        JSONObject topObject = new JSONObject();
        for (Object object : array) {
            JSONObject item = (JSONObject) object;
            String id = item.getString(idKey);
            String pid = item.getString(pidKey);
            if(!id.equals(pid)) {
                tempArray.add(item);
            } else {
                topObject = item;
            }
        }

        JSONArray sortedArray = new JSONArray();
        JSONArray subArray = getChildTree(tempArray, idKey, pidKey, childKey, topName);
        topObject.put(childKey, subArray);
        sortedArray.add(topObject);
        return sortedArray;
    } else {
        return getChildTree(array, idKey, pidKey, childKey, topName);
    }
}

/**
 * 判断是否有顶级节点
 * @param array
 * @param idKey
 * @param pidKey
 * @param topName
 * @return
 */
public static boolean hasTopNode(JSONArray array, String idKey, String pidKey, String topName) {
    for (Object object : array) {
        JSONObject item = (JSONObject) object;
        String id = item.getString(idKey);
        String pid = item.getString(pidKey);
        if(topName.equals(pid) && id.equals(pid)) {
            return true;
        }
    }
    return false;
}

/**
 * 获取子树
 * @param array
 * @param idKey
 * @param pidKey
 * @param childKey
 * @param topName
 * @return
 */
private static JSONArray getChildTree(JSONArray array, String idKey, String pidKey, String childKey, String topName) {
    JSONArray sortedArray = new JSONArray();
    for (Object object : array) {
        JSONObject item = (JSONObject) object;
        String id = item.getString(idKey);
        String pid = item.getString(pidKey);
        if(topName.equals(pid)) {
            if(getChildTree(array, idKey, pidKey, childKey, id).size() > 0) {
                JSONArray subArray = getChildTree(array, idKey, pidKey, childKey, id);
                item.put(childKey, subArray);
            }
            sortedArray.add(item);
        }
    }
    return sortedArray;
}