字节跳动（社招）一面算法原题

宫水三叶的刷题日记

发布于 2024-05-13 16:45:13

1710

发布于 2024-05-13 16:45:13

文章被收录于专栏：宫水三叶的刷题日记

题目描述

平台：LeetCode

题号：1129

在一个有向图中，节点分别标记为 0, 1, ..., n-1。

图中每条边为红色或者蓝色，且存在自环或平行边。

red_edges 中的每一个 [i, j] 对表示从节点 i 到节点 j 的红色有向边。

类似地，blue_edges 中的每一个 [i, j] 对表示从节点 i 到节点 j 的蓝色有向边。

返回长度为 n 的数组 answer，其中 answer[X] 是从节点 0 到节点 X 的红色边和蓝色边交替出现的最短路径的长度。

如果不存在这样的路径，那么 answer[x] = -1。

示例 1：

输入：n = 3, red_edges = [[0,1],[1,2]], blue_edges = []

输出：[0,1,-1]

示例 2：

输入：n = 3, red_edges = [[0,1]], blue_edges = [[2,1]]

输出：[0,1,-1]

示例 3：

输入：n = 3, red_edges = [[1,0]], blue_edges = [[2,1]]

输出：[0,-1,-1]

示例 4：

输入：n = 3, red_edges = [[0,1]], blue_edges = [[1,2]]

输出：[0,1,2]

示例 5：

输入：n = 3, red_edges = [[0,1],[0,2]], blue_edges = [[1,0]]

输出：[0,1,1]

提示：

1 <= n <= 100
red_edges.length <= 400
blue_edges.length <= 400
red_edges[i].length = blue_edges[i].length = 2
0 <= red_edges[i][j], blue_edges[i][j] < n

朴素 BFS

为了方便，将 redEdges 记为 rs，将 blueEdges 记为 bs。

使用两数组进行建图，利用点数和边数关系（稀疏图），使用「邻接表」方式进行建图。

将所有红色有向边权值记为 1，将所有蓝色有向边权值记为 -1。注意这里的权重仅表示该边的颜色，并非代表经过该边的真实成本。

也正是经过所有边的成本均相同，对于原问题「从 0 号节点进行出发，求到所有点的颜色交替的最短路径」，我们容易想到使用 BFS 进行求解。

BFS 过程中将三元组

(point, last, step)

进行入队：

point 代表当前所在点编号
last 代表到达 point 所使用的边的颜色，默认为 0 代表不需要经过任何边到达起点
step 代表到达 point 所使用的步数

利用数据范围不大（点数为

100

，边数为

800

），可以直接对每个点进行一次独立的 BFS 来求最短路。由于图存在重边和自环，因此在求解最短路的过程需要对边进行去重。

Java 代码：

class Solution {
    static int N = 110, M = 810, idx = 0;
    static int[] he = new int[N], e = new int[M], ne = new int[M], w = new int[M];
    void add(int a, int b, int c) {
        e[idx] = b;
        w[idx] = c;
        ne[idx] = he[a];
        he[a] = idx++;
    }
    public int[] shortestAlternatingPaths(int n, int[][] rs, int[][] bs) {
        idx = 0;
        Arrays.fill(he, -1);
        for (int[] e : rs) add(e[0], e[1], 1);
        for (int[] e : bs) add(e[0], e[1], -1);
        int[] ans = new int[n];
        boolean[] vis = new boolean[rs.length + bs.length];
        out:for (int k = 1; k < n; k++) {
            ans[k] = -1;
            Arrays.fill(vis, false);
            Deque<int[]> d = new ArrayDeque<>();
            d.addLast(new int[]{0, 0, 0}); // point, last, step
            while (!d.isEmpty()) {
                int[] info = d.pollFirst();
                int p = info[0], last = info[1], step = info[2];
                for (int i = he[p]; i != -1; i = ne[i]) {
                    int j = e[i], c = w[i];
                    if (vis[i]) continue;
                    if (c + last == 0 || last == 0) {
                        if (j == k) {
                            ans[k] = step + 1;
                            continue out;
                        } else {
                            d.addLast(new int[]{j, c, step + 1});
                            vis[i] = true;
                        }
                    }
                }
            }
        }
        return ans;
    }
}

C++ 代码：

class Solution {
public:
    int he[110], e[810], ne[810], w[810], idx = 0;
    void add(int a, int b, int c) {
        e[idx] = b;
        w[idx] = c;
        ne[idx] = he[a];
        he[a] = idx++;
    }
    vector<int> shortestAlternatingPaths(int n, vector<vector<int>>& rs, vector<vector<int>>& bs) {
        idx = 0;
        fill(he, he + n, -1);
        for (auto& e : rs) add(e[0], e[1], 1);
        for (auto& e : bs) add(e[0], e[1], -1);
        vector<int> ans(n);
        vector<bool> vis(rs.size() + bs.size());
        for (int k = 1; k < n; k++) {
            ans[k] = -1;
            fill(vis.begin(), vis.end(), false);
            queue<vector<int>> q;
            q.push({ 0, 0, 0 }); // point, last, step
            bool out = false;
            while (!q.empty()) {
                if (out) break;
                vector<int> info = q.front(); q.pop();
                int p = info[0], last = info[1], step = info[2];
                for (int i = he[p]; i != -1; i = ne[i]) {
                    int j = e[i], c = w[i];
                    if (vis[i]) continue;
                    if (c + last == 0 || last == 0) {
                        if (j == k) {
                            ans[k] = step + 1;
                            out = true;
                            break;
                        } else {
                            q.push({ j, c, step + 1 });
                            vis[i] = true;
                        }
                    }
                }
            }
        }
        return ans;
    }
};

时间复杂度：将点数记为 n ，边数记为 m。建图复杂度为

O(m)

；构造答案复杂度为

O(n \times (n + m))

。整体复杂度为

O(n \times (n + m))

空间复杂度：

O(n + m)

优化

实际上，我们并没有对每个点进行独立 BFS 的必要。

为了获取所有从节点 0 出发的最短路，可直接从节点 0 进行出发（对边进行去重），所有能够从节点 0 沿交替路径到达的节点必然都会被访问到，且节点首次被访问到时必然是最短路径。

因此我们只需要初始化所有的 ans[i] = -1，并从节点 0 开始进行一次 BFS 即可。

Java 代码：

class Solution {
    static int N = 110, M = 810, idx = 0;
    static int[] he = new int[N], e = new int[M], ne = new int[M], w = new int[M];
    void add(int a, int b, int c) {
        e[idx] = b;
        w[idx] = c;
        ne[idx] = he[a];
        he[a] = idx++;
    }
    public int[] shortestAlternatingPaths(int n, int[][] rs, int[][] bs) {
        idx = 0;
        Arrays.fill(he, -1);
        for (int[] e : rs) add(e[0], e[1], 1);
        for (int[] e : bs) add(e[0], e[1], -1);
        int[] ans = new int[n];
        Arrays.fill(ans, -1);
        ans[0] = 0;
        boolean[] vis = new boolean[rs.length + bs.length];
        Arrays.fill(vis, false);
        Deque<int[]> d = new ArrayDeque<>();
        d.addLast(new int[]{0, 0, 0}); // point, last, step
        while (!d.isEmpty()) {
            int[] info = d.pollFirst();
            int p = info[0], last = info[1], step = info[2];
            for (int i = he[p]; i != -1; i = ne[i]) {
                if (vis[i]) continue;
                int j = e[i], c = w[i];
                if (c + last == 0 || last == 0) {
                    ans[j] = ans[j] == -1 ? step + 1 : Math.min(ans[j], step + 1);
                    d.addLast(new int[]{j, c, step + 1});
                    vis[i] = true;
                }
            }
        }
        return ans;
    }
}

C++ 代码：

class Solution {
public:
    int he[110], e[810], ne[810], w[810], idx = 0;
    void add(int a, int b, int c) {
        e[idx] = b;
        w[idx] = c;
        ne[idx] = he[a];
        he[a] = idx++;
    }
    vector<int> shortestAlternatingPaths(int n, vector<vector<int>>& rs, vector<vector<int>>& bs) {
        idx = 0;
        fill(he, he + n, -1);
        for (auto& e : rs) add(e[0], e[1], 1);
        for (auto& e : bs) add(e[0], e[1], -1);
        vector<int> ans(n, -1);
        ans[0] = 0;
        vector<bool> vis(rs.size() + bs.size(), false);
        deque<vector<int>> d;
        d.push_back({0, 0, 0}); // point, last, step
        while (!d.empty()) {
            vector<int> info = d.front();
            d.pop_front();
            int p = info[0], last = info[1], step = info[2];
            for (int i = he[p]; i != -1; i = ne[i]) {
                if (vis[i]) continue;
                int j = e[i], c = w[i];
                if (c + last == 0 || last == 0) {
                    ans[j] = ans[j] == -1 ? step + 1 : min(ans[j], step + 1);
                    d.push_back({j, c, step + 1});
                    vis[i] = true;
                }
            }
        }
        return ans;
    }
};