package com.zy.core.utils;

import java.util.*;

public class TrafficControlUtils {

    public static List<List<String>> groupNodes(Collection<String> keys) {
        // 1. 按位置聚类：Map<位置键, 节点列表>
        Map<String, List<String>> clusterMap = new HashMap<>();
        for (String key : keys) {
            String rowStr = key.substring(0, 2);
            String colStr = key.substring(2, 5);
            String posKey = rowStr + "_" + colStr; // 位置键格式: "RR_CCC"
            clusterMap.computeIfAbsent(posKey, k -> new ArrayList<>()).add(key);
        }

        // 2. 初始化并查集
        Map<String, String> parentMap = new HashMap<>();
        Map<String, Integer> rankMap = new HashMap<>();
        for (String posKey : clusterMap.keySet()) {
            parentMap.put(posKey, posKey); // 初始父节点指向自己
            rankMap.put(posKey, 0);
        }

        // 3. 遍历所有位置键，合并相邻簇
        for (String posKey : clusterMap.keySet()) {
            String[] parts = posKey.split("_");
            int row = Integer.parseInt(parts[0]);
            int col = Integer.parseInt(parts[1]);

            // 检查四个方向：左、右、上、下
            int[][] directions = {{-1, 0}, {1, 0}, {0, -1}, {0, 1}}; // {dRow, dCol}
            for (int[] dir : directions) {
                int newRow = row + dir[0];
                int newCol = col + dir[1];
                String neighborKey = String.format("%02d_%03d", newRow, newCol);

                if (parentMap.containsKey(neighborKey)) {
                    union(posKey, neighborKey, parentMap, rankMap);
                }
            }
        }

        // 4. 生成分组结果
        Map<String, List<String>> groupMap = new HashMap<>();
        for (String posKey : clusterMap.keySet()) {
            String root = find(posKey, parentMap);
            groupMap.computeIfAbsent(root, k -> new ArrayList<>())
                    .addAll(clusterMap.get(posKey));
        }

        return new ArrayList<>(groupMap.values());
    }

    // 并查集：查找根节点（带路径压缩）
    private static String find(String x, Map<String, String> parentMap) {
        if (!parentMap.get(x).equals(x)) {
            parentMap.put(x, find(parentMap.get(x), parentMap));
        }
        return parentMap.get(x);
    }

    // 并查集：按秩合并
    private static void union(String x, String y, Map<String, String> parentMap, Map<String, Integer> rankMap) {
        String rootX = find(x, parentMap);
        String rootY = find(y, parentMap);
        if (rootX.equals(rootY)) return;

        int rankX = rankMap.get(rootX);
        int rankY = rankMap.get(rootY);
        if (rankX < rankY) {
            parentMap.put(rootX, rootY);
        } else if (rankX > rankY) {
            parentMap.put(rootY, rootX);
        } else {
            parentMap.put(rootY, rootX);
            rankMap.put(rootX, rankX + 1);
        }
    }

}