3559. Number of Ways to Assign Edge Weights II

#include <stdlib.h>

// 演算法 / Algorithm:
// 1) 觀察:路徑長度 k 條邊時,使成本為奇數的方法數 = 2^(k-1)(k=0 時為 0)。
//    For a path with k edges, the number of odd-cost assignments is 2^(k-1) (0 if k==0).
// 2) 用 BFS 求深度與父節點,倍增建表求 LCA,得到兩點距離 k。
//    BFS for depth/parent, binary-lifting LCA to get distance k per query.

#define MOD 1000000007        // 模數,質數 / the prime modulus
#define LOG 17                // 2^17 = 131072 > 1e5,夠跳到根 / enough jump levels for n up to 1e5

int* assignEdgeWeights(int** edges, int edgesSize, int* edgesColSize,
                       int** queries, int queriesSize, int* queriesColSize,
                       int* returnSize) {
    int n = edgesSize + 1;                       // 節點數 = 邊數 + 1 / nodes = edges + 1

    // ---- 建鄰接表(用「頭+next」的陣列鏈結,避免動態 vector)----
    // Build adjacency list with head/next arrays (linked list inside arrays).
    int *head = malloc((n + 1) * sizeof(int));   // head[u] = u 的第一條邊索引 / first edge index of u
    int *nxt  = malloc(2 * edgesSize * sizeof(int)); // 無向邊存兩次 / each undirected edge stored twice
    int *to   = malloc(2 * edgesSize * sizeof(int)); // to[e] = 邊 e 的另一端 / the other endpoint of edge e
    for (int i = 1; i <= n; i++) head[i] = -1;   // -1 代表「沒有邊」/ -1 means "no edge"
    int cnt = 0;                                  // 已加入的有向邊數 / number of directed edges added
    for (int i = 0; i < edgesSize; i++) {
        int u = edges[i][0], v = edges[i][1];    // 取出這條邊的兩端 / the two endpoints
        to[cnt] = v; nxt[cnt] = head[u]; head[u] = cnt++; // 把 u->v 插到 u 的鏈表頭 / push u->v
        to[cnt] = u; nxt[cnt] = head[v]; head[v] = cnt++; // 把 v->u 插到 v 的鏈表頭 / push v->u
    }

    // ---- BFS 從根 1 算 depth 和 up[0](直接父節點)----
    // BFS from root 1 to fill depth and up[0] (immediate parent).
    int *depth = malloc((n + 1) * sizeof(int));  // depth[v] = v 到根的邊數 / edges from v to root
    int **up = malloc(LOG * sizeof(int*));        // up[k][v] = v 往上 2^k 步的祖先 / 2^k-th ancestor
    for (int k = 0; k < LOG; k++) up[k] = malloc((n + 1) * sizeof(int));
    int *queue = malloc((n + 1) * sizeof(int));  // BFS 用的佇列 / the BFS queue
    char *vis = calloc(n + 1, 1);                 // 是否已訪問,calloc 全設 0 / visited flags, zeroed
    int qh = 0, qt = 0;                           // 佇列頭尾指標 / queue head & tail
    queue[qt++] = 1; vis[1] = 1; depth[1] = 0; up[0][1] = 0; // 根的父節點用 0 當哨兵 / root parent = 0 sentinel
    while (qh < qt) {                             // 佇列非空就繼續 / while queue not empty
        int u = queue[qh++];                      // 取出隊首 / pop front
        for (int e = head[u]; e != -1; e = nxt[e]) { // 走訪 u 的每條邊 / iterate u's edges
            int w = to[e];                        // 鄰居 / neighbor
            if (!vis[w]) {                        // 第一次看到才處理 / only if unvisited
                vis[w] = 1;
                depth[w] = depth[u] + 1;          // 子節點深度 = 父深度 + 1 / child depth
                up[0][w] = u;                     // w 的直接父節點是 u / parent of w is u
                queue[qt++] = w;                  // 入列 / enqueue
            }
        }
    }
    depth[0] = 0;                                 // 哨兵節點 0 的深度 / sentinel depth

    // ---- 用倍增填出整張祖先表 / Fill the full ancestor table ----
    for (int k = 1; k < LOG; k++) {
        up[k][0] = 0;                             // 哨兵往上還是哨兵 / sentinel stays sentinel
        for (int v = 1; v <= n; v++)
            // 往上 2^k 步 = 先跳 2^(k-1),再跳 2^(k-1) / a 2^k jump is two 2^(k-1) jumps
            up[k][v] = up[k - 1][ up[k - 1][v] ];
    }

    // ---- 預先算 2 的各次方 / Precompute powers of two ----
    long long *pw = malloc((n + 1) * sizeof(long long));
    pw[0] = 1;                                    // 2^0 = 1
    for (int i = 1; i <= n; i++)
        pw[i] = pw[i - 1] * 2 % MOD;              // 每步取模避免溢位 / mod each step to avoid overflow

    // ---- 逐一回答查詢 / Answer each query ----
    int* ans = malloc(queriesSize * sizeof(int));
    for (int i = 0; i < queriesSize; i++) {
        int a = queries[i][0], b = queries[i][1]; // 這次查詢的兩點 / the two query nodes
        int u = a, v = b;                         // 留著原始值算距離 / keep originals for distance
        if (depth[a] < depth[b]) { int t = a; a = b; b = t; } // 讓 a 比較深 / make a the deeper one
        int diff = depth[a] - depth[b];           // 深度差 / depth difference
        for (int k = 0; k < LOG; k++)             // 把 a 往上拉到和 b 同深度 / lift a up to b's depth
            if (diff & (1 << k)) a = up[k][a];    // diff 的第 k 個位元是 1 就跳 2^k / jump 2^k if bit k set
        int lca;
        if (a == b) {
            lca = a;                              // b 本來就是 a 的祖先 / b was already an ancestor of a
        } else {
            for (int k = LOG - 1; k >= 0; k--)    // 兩點一起往上跳到「父節點剛好相同」/ lift both together
                if (up[k][a] != up[k][b]) { a = up[k][a]; b = up[k][b]; }
            lca = up[0][a];                       // 再上一步就是 LCA / one more step is the LCA
        }
        int dist = depth[u] + depth[v] - 2 * depth[lca]; // 路徑邊數 / number of edges on the path
        ans[i] = dist == 0 ? 0 : (int)pw[dist - 1];      // k=0 答 0,否則 2^(k-1) / 0 if no edges else 2^(k-1)
    }

    *returnSize = queriesSize;                    // 回傳陣列長度 / set output length
    // 釋放暫時用的記憶體(ans 要留給呼叫者)/ free scratch memory (ans is returned to caller)
    free(head); free(nxt); free(to); free(depth);
    for (int k = 0; k < LOG; k++) free(up[k]);
    free(up); free(queue); free(vis); free(pw);
    return ans;
}

#include <vector>
using namespace std;

// 演算法 / Algorithm:
// 路徑有 k 條邊時,使成本為奇數的方法數 = 2^(k-1)(k=0 時為 0)。
// With k edges on the path, odd-cost assignments = 2^(k-1) (0 when k==0).
// 用 BFS 求深度/父節點,倍增 LCA 求兩點距離 k。
// BFS for depth/parent, binary-lifting LCA for the distance k.

class Solution {
public:
    static const int MOD = 1000000007;            // 模數 / the modulus
    static const int LOG = 17;                     // 跳躍層數,2^17 > 1e5 / lifting levels

    vector<int> assignEdgeWeights(vector<vector<int>>& edges,
                                  vector<vector<int>>& queries) {
        int n = edges.size() + 1;                  // 節點數 / number of nodes

        // 鄰接表:adj[u] 是 u 所有鄰居的清單 / adjacency list: adj[u] holds u's neighbors
        vector<vector<int>> adj(n + 1);
        for (auto& e : edges) {                    // range-for 走訪每條邊 / range-based for over edges
            adj[e[0]].push_back(e[1]);             // 無向邊兩個方向都加 / add both directions
            adj[e[1]].push_back(e[0]);
        }

        vector<int> depth(n + 1, 0);               // 每點深度 / depth of each node
        // up 是二維表,up[k][v] = v 往上 2^k 步的祖先 / up[k][v] = 2^k-th ancestor of v
        vector<vector<int>> up(LOG, vector<int>(n + 1, 0));

        // BFS 從根 1 / BFS from root 1
        vector<char> vis(n + 1, 0);                // 訪問標記 / visited flags
        vector<int> q;                             // 用 vector 當佇列 / vector used as a queue
        q.push_back(1); vis[1] = 1; up[0][1] = 0;  // 根的父節點是哨兵 0 / root's parent = sentinel 0
        for (size_t i = 0; i < q.size(); i++) {    // i 隨 q 變長持續推進 / i advances as q grows
            int u = q[i];                          // 取出目前節點 / current node
            for (int w : adj[u]) {                 // 走訪鄰居 / iterate neighbors
                if (!vis[w]) {                     // 沒訪問過才展開 / expand only if unvisited
                    vis[w] = 1;
                    depth[w] = depth[u] + 1;       // 子節點深度 / child depth
                    up[0][w] = u;                  // 直接父節點 / immediate parent
                    q.push_back(w);                // 入列 / enqueue
                }
            }
        }

        // 倍增填表 / build the lifting table
        for (int k = 1; k < LOG; k++)
            for (int v = 1; v <= n; v++)
                // 2^k 步 = 兩次 2^(k-1) 步 / a 2^k jump = two 2^(k-1) jumps
                up[k][v] = up[k - 1][up[k - 1][v]];

        // 預先算 2 的次方 / precompute powers of two
        vector<long long> pw(n + 1);
        pw[0] = 1;
        for (int i = 1; i <= n; i++) pw[i] = pw[i - 1] * 2 % MOD;

        // 求 LCA 的 lambda / a lambda computing the LCA of two nodes
        auto lca = [&](int a, int b) {
            if (depth[a] < depth[b]) swap(a, b);   // 讓 a 較深 / make a the deeper one
            int diff = depth[a] - depth[b];        // 深度差 / depth gap
            for (int k = 0; k < LOG; k++)          // 把 a 拉到和 b 同深度 / lift a to b's depth
                if (diff & (1 << k)) a = up[k][a]; // 第 k 位元為 1 就跳 2^k / jump 2^k if bit set
            if (a == b) return a;                  // b 已是祖先 / b is already the ancestor
            for (int k = LOG - 1; k >= 0; k--)     // 兩點一起往上 / lift both together
                if (up[k][a] != up[k][b]) { a = up[k][a]; b = up[k][b]; }
            return up[0][a];                       // 再一步即 LCA / one step up is the LCA
        };

        vector<int> ans;
        ans.reserve(queries.size());               // 預留空間避免重複擴容 / reserve to avoid reallocs
        for (auto& qy : queries) {
            int u = qy[0], v = qy[1];              // 結構化取兩端 / the two query endpoints
            int l = lca(u, v);                     // 最近公共祖先 / their LCA
            int dist = depth[u] + depth[v] - 2 * depth[l]; // 路徑邊數 / edges on the path
            ans.push_back(dist == 0 ? 0 : (int)pw[dist - 1]); // 0 或 2^(k-1) / 0 or 2^(k-1)
        }
        return ans;
    }
};