• 难度：中等
• 标签：优先队列、二叉查找树、深度优先搜索
• Amazon 面试题

题目

给一些 movie（编号从 0 开始）的 rating 和它们的联系关系，联系关系可以传递（a 和 b 有联系，b 和 c 有联系，a 和 c 也认为有联系）。给出每个 movie 的直接联系 list。再给定一个 movie 的编号为 S，找到和 S 相联系的 movie 中的 rating 最高的 K 个 movie（当与 S 相联系的 movie 少于 K 个时，输出所有。输出答案时可以按照任意次序，注意联系 movie 并不包括 S。）

注意事项

• movie 个数 n ≤ 20000
• 保证编号为 0~n-1（n 为 movie 个数）
• 保证联系的边，2 个编号都属于 0 ~ n-1。
• 保证总边数 m ≤ 100000
• 保证每个 movie 的 rating 都为 int 范围内的整数，且互不相同。

样例

Sample Input:
[10, 20, 30, 40]
[
  [1, 3],
  [0, 2],
  [1],
  [0]
]
0
2

Sample Output:
[2, 3]

解释：
在 contactRelationship 中，0 号与 [1, 3] 有联系，1 号与 [0, 2] 有联系，2 号与 [1] 有联系，3 号与 [0] 有联系
所以最终输出 0 号和 movie 相联系的有 [1, 2, 3]，按它们的 rating 从高到低排列为 [3, 2, 1]，所以输出 [2, 3]。

Sample Input:
[10, 20, 30, 40, 50, 60, 70, 80, 90]
[
  [1, 4, 5],
  [0, 2, 3],
  [1, 7],
  [1, 6, 7],
  [0],
  [0],
  [3],
  [2, 3],
  []
]
5
3

Sample Output:
[6, 7, 4]

解释：
在 contactRelationship 中，0 号与 [1, 4, 5] 有联系，1 号与 [0, 2, 3] 有联系，2 号与 [1, 7] 有联系，3 号与 [1, 6, 7] 有联系，4 号与 [0] 有联系，5 号与 [0] 有联系，6 号与 [3] 有联系，7 号与 [2, 3] 有联系，8 号无任何联系。
最终和 5 号 movie 相联系的有 [0, 1, 2, 3, 4, 6, 7] 按它们的 rating 从高到低排列为 [7, 6, 4, 3, 2, 1, 0]，注意 8 号 movie 不和 5 号 movie 相连，所以它 rating 最高但是不计入答案。

解答

class Solution 
{
public:
    /**
     * @param rating: the rating of the movies
     * @param G: the relationship of movies
     * @param S: the begin movie
     * @param K: top K rating
     * @return : the top k largest rating movie which contact with S
     */
    vector<int> topKMovie(vector<int> &rating, vector<vector<int>> &G, int S, int k) 
    {
        vector<bool> isJoined(rating.size(), false);    // 设置标记，避免元素重复入队
        queue<int> movieQueue;  // 设置队列，用于寻找关系
        priority_queue<pair<int, int>> ratingHeap;  // 优先队列，存放分数及对应的电影编号
        vector<int> result;
        isJoined[S] = true; // 将起始电影标记为 true
        for(auto x : G[S])  // 把与起始电影有关系的电影入队
            movieQueue.push(x);
        while(!movieQueue.empty()) 
        {
            int temp = movieQueue.front();
            movieQueue.pop();
            if(isJoined[temp])
                continue;
            isJoined[temp] = true;
            ratingHeap.push(pair<int, int>(rating[temp], temp));
            for(auto y : G[temp]) 
            {
                if(!isJoined[y])
                    movieQueue.push(y);
            }
        }
        // 输出结果
        while(k--) 
        {
            result.push_back(ratingHeap.top().second);
            ratingHeap.pop();
        }
        return result;
    }
};