美团笔试（2020-9-6）员工分配问题

算法 美团笔试（2020-9-6）员工分配问题

@author:Jingdai
@date:2020.11.18

前几天看美团笔试题，记录一下。该题是美团9月6号的算法笔试题第4题。

题目描述

公司有 n 位员工，需要划分 n 位员工的从属关系，划分要求如下：

1. 每个人要么没有下属，要么至少有 2 个直接下属
2. 第 i 个人的下属（包括自己）有 ai 个

注意：直接下属和下属（包括自己）可以分别看成该员工的“儿子”和“子树”，问是否存在这样一种关系。

输入描述：

• 输入第一行是一个整数 n （n<=24），表示公司有 n 个人。

• 接下来一行 n 个数，第 i 个数为 ai

思路

这个题显然是一个树的问题，先分析一下具体的题意。如图，先看几个满足题目要求的例子再看，能更好的理解题意。

对于要求1，其实意思就是对于树中的每个非叶子节点，要求其儿子数大于等于2。也就是树中的每个节点的度都不能等于1。

对于要求2，意思是每个节点值（ai）代表了该节点的子树（包括自己）中有 ai 个节点。那么对于这个树的根节点，就代表了整个树的节点数，所以最大的 ai 是根节点，同时若能构成这样的树，最大的 ai 一定等于 n。所以可以根据最大的 ai 是否等于 n 进行一个预处理判断。

同时结合要求1和要求2，会发现 ai 值如果等于2，那一定不能构成这样的树，因为 ai 最小为1（根据要求2），而一个非叶子节点要至少包含 2 个叶子节点，那么非叶子节点值最小为3，所以同样可以根据输入值中是否包含2 来进行预处理判断。

接下来看具体怎么解这个问题，这里利用回溯算法进行求解。首先看变量的定义：

• nodes 数组

  nodes 数组记录输入的 ai 值。首先将 nodes 从大到小排序，那么 nodes[0] 就是树的根节点，除了 nodes[0] 都是有父节点的。这里将 nodes 数组中的值看成还需要分配的节点数，比如 nodes[i] = 6 代表以该节点为根的子树还需要分配 5 个（自己算1个）节点；当 nodes[i] = 1 代表该节点不需要再分配子节点了。

• children 数组

  记录每个节点的孩子数，用于最后判断是否满足每个节点的度都不为1。

• unfinishedParentSet 集合

  这个 set 集合代表还需要分配子节点的父节点下标，即 nodes[i] > 1 的所有节点。当这个 set 为空代表已经没有需要分配的父节点了。

接下来看算法。首先预处理，如前所述，对于每个输入，如果 ai 等于 2 ，代表不能构成这样的树，输入完成后，对 nodes 进行从大到小排序，如果 nodes[0] 不等于 n，代表不能构成这样的树。

预处理完成后，对于不能判断的，进行下一步。首先遍历 nodes 数组，将 nodes 数组中大于1的下标加入unfinishedParentSet。然后进行深度优先遍历并回溯。

如图，将 nodes 数组元素从 nodes[1]（nodes[0] 是根节点，没有父节点）开始尝试去分配给 unfinishedParentSet 里的父节点，图中的父x节点都是 unfinishedParentSet 中的节点，一个父节点分配失败就尝试下一个父节点，有成功的就返回 true，全部父节点都不行就返回 false。

尝试分配即 nodes[i] -= nodes[child]，将 children[i] 加1，同时如果 nodes[i] 等于1后代表该父节点分配完毕，从 unfinishedParentSet 中删去。 如果尝试分配失败就回溯，将 nodes 、children 和 unfinishedParentSet 还原。同时，由于题目说明节点度不能为1，可以根据此进行剪枝，当 nodes[i] - nodes[child] == 1 && children[i] = 0 时，不进行分配，因为这样分配会使节点 i 的度为1，就不用往下走了，进行剪枝。

同时注意一个细节，按算法逻辑来说从头到尾用一个 unfinishedParentSet 就行，但是我们在遍历过程中会对 set 集合中的元素进行增减，Java的 foreach 遍历中不能增减元素（会抛异常），iterator 遍历中不能增加元素，普通的for遍历增减元素会影响遍历的语意（删除一个元素可能会少遍历到元素），所以代码中每一次分配时都会重新构建一个和原来一样的 set ，对新的 set 进行增减。当然你也可以多用一个字段记录该节点是否被删除。这里和具体的语言实现有关，和算法没什么关系。

具体代码如下。

代码

import java.util.*;

public class Solution {

    // number of nodes
    public static int n;

    public static int[] nodes;

    // the number of the node i's children
    public static int[] children;

    public static void main(String[] args) {
        
        Scanner in = new Scanner(System.in);
        
        n = in.nextInt();
        nodes = new int[n];
        children = new int[n];

        in.nextLine();
        
        boolean canMakeTree = true;

        for (int i = 0; i < n; i++) {
            nodes[i] = in.nextInt();
            if (nodes[i] == 2)
                canMakeTree = false;
        }

        Arrays.sort(nodes);

        // reverse
        for (int i = 0; i < n/2; i++) {
            int temp = nodes[i];
            nodes[i] = nodes[n-1-i];
            nodes[n-1-i] = temp;
        }

        if (nodes[0] != n) {
            canMakeTree = false;
        }

        if (!canMakeTree) {
            System.out.println("NO");
            return;
        }

        Set<Integer> unfinishedParentSet = new HashSet<>();
        
        for (int i = 0; i < n; i++) {
            if (nodes[i] > 1)
                unfinishedParentSet.add(i);
        }

        if (allocateNode(1, unfinishedParentSet)) {
            System.out.println("YES");
        } else {
            System.out.println("NO");
        }
    }

    public static boolean allocateNode(int index, Set<Integer> unfinishedParentSet) {
        
        Set<Integer> newSet = new HashSet<>(unfinishedParentSet);

        if (index == n) {
            if (unfinishedParentSet.size() != 0)
                return false;
            for (int i : children) {
                if (i == 1)
                    return false;
            }  
            return true;
        }

        for (int parentIndex : unfinishedParentSet) {
            if (nodes[parentIndex] > nodes[index]) {
                if (nodes[parentIndex] == nodes[index] + 1 
                    && children[parentIndex] == 0) {
                    continue;
                }
                // try to allocate
                nodes[parentIndex] -= nodes[index];
                if (nodes[parentIndex] == 1)
                    newSet.remove(parentIndex);
                children[parentIndex]++;

                // succeed
                if (allocateNode(index + 1, newSet))
                    return true;
                
                // fail to allocate
                if (nodes[parentIndex] == 1)
                    newSet.add(parentIndex);
                nodes[parentIndex] += nodes[index];
                children[parentIndex]--;
            }
        }
        return false;
    }
}

参考

• 题目参考
• 解法思路