Kimi LeetCode 3515. 带权树中的最短路径 C语言实现
这是 LeetCode 3515「带权树中的最短路径」的 C 语言实现,采用 DFS 欧拉序 + 线段树(区间加、单点查) 的经典做法,并使用迭代 DFS 避免链状树递归栈溢出。
---
核心思路
1. 以 1 为根 DFS,记录每个节点的进入时间 `tin[i]`、离开时间 `tout[i]` 和到根的初始距离 `dist[i]`。
2. 子树 = 连续区间:节点 `i` 的子树中所有节点的时间戳落在 `[tin[i], tout[i]]`。
3. 修改边权 `(u, v)`:`tin` 较大的节点是子节点 `child`,该边权变化 `delta` 会影响 `child` 整个子树到根的距离,即对区间 `[tin[child], tout[child]]` 整体加 `delta`。
4. 线段树:维护区间加、单点查。查询节点 `x` 时,返回位置 `tin[x]` 的值即可。
---
C 代码
```c
#include <stdlib.h>
/* ============ 线段树:区间加、单点查(懒标记不下传,查询时累加) ============ */
typedef struct {
long long *lazy;
int n;
} SegTree;
static SegTree* segCreate(int n) {
SegTree *s = malloc(sizeof(SegTree));
s->n = n;
s->lazy = calloc(4 * n + 4, sizeof(long long));
return s;
}
static void segFree(SegTree *s) {
free(s->lazy);
free(s);
}
/* 区间 [ql,qr] 增加 val */
static void segAdd(SegTree *s, int idx, int l, int r, int ql, int qr, long long val) {
if (qr < l || r < ql) return;
if (ql <= l && r <= qr) {
s->lazy[idx] += val;
return;
}
int mid = (l + r) >> 1;
segAdd(s, idx << 1, l, mid, ql, qr, val);
segAdd(s, idx << 1 | 1, mid + 1, r, ql, qr, val);
}
/* 单点查询 */
static long long segQuery(SegTree *s, int idx, int l, int r, int pos) {
long long res = s->lazy[idx];
if (l == r) return res;
int mid = (l + r) >> 1;
if (pos <= mid) res += segQuery(s, idx << 1, l, mid, pos);
else res += segQuery(s, idx << 1 | 1, mid + 1, r, pos);
return res;
}
/* ============ 邻接表 ============ */
typedef struct Edge {
int to, w;
struct Edge *next;
} Edge;
/* ============ 全局 DFS 变量 ============ */
static int *tin, *tout, timestamp;
static long long *dist;
/* 迭代 DFS(防止链状树 1e5 层递归栈溢出) */
static void dfs(Edge **g, int n) {
typedef struct {
int u, fa;
long long d;
Edge *next;
int state; /* 0=首次进入, 1=处理完子节点 */
} Stk;
Stk *stk = malloc((n + 1) * sizeof(Stk));
int top = 0;
stk[top++] = (Stk){1, 0, 0, g[1], 0};
while (top > 0) {
Stk *cur = &stk[top - 1];
int u = cur->u;
if (cur->state == 0) {
tin[u] = ++timestamp;
dist[u] = cur->d;
cur->state = 1;
}
/* 跳过指向父节点的边 */
while (cur->next && cur->next->to == cur->fa)
cur->next = cur->next->next;
if (cur->next) {
Edge *e = cur->next;
cur->next = e->next;
stk[top++] = (Stk){e->to, u, cur->d + e->w, g[e->to], 0};
} else {
tout[u] = timestamp;
top--;
}
}
free(stk);
}
/* ============ LeetCode 接口 ============ */
int* treeQueries(int n, int** edges, int edgesSize, int* edgesColSize,
int** queries, int queriesSize, int* queriesColSize,
int* returnSize) {
tin = malloc((n + 1) * sizeof(int));
tout = malloc((n + 1) * sizeof(int));
dist = malloc((n + 1) * sizeof(long long));
Edge **g = calloc(n + 1, sizeof(Edge*));
/* 建无向图 */
for (int i = 0; i < edgesSize; i++) {
int u = edges[i][0], v = edges[i][1], w = edges[i][2];
Edge *e1 = malloc(sizeof(Edge));
e1->to = v; e1->w = w; e1->next = g[u]; g[u] = e1;
Edge *e2 = malloc(sizeof(Edge));
e2->to = u; e2->w = w; e2->next = g[v]; g[v] = e2;
}
/* 求 tin / tout / dist */
timestamp = 0;
dfs(g, n);
/* 线段树:单点初始化 dist */
SegTree *seg = segCreate(n);
for (int i = 1; i <= n; i++)
segAdd(seg, 1, 1, n, tin[i], tin[i], dist[i]);
int *ans = malloc(queriesSize * sizeof(int));
int cnt = 0;
for (int i = 0; i < queriesSize; i++) {
if (queries[i][0] == 1) { /* 修改边权 */
int u = queries[i][1], v = queries[i][2], newW = queries[i][3];
int oldW = 0;
/* 读取旧权值并更新邻接表 */
for (Edge *e = g[u]; e; e = e->next)
if (e->to == v) { oldW = e->w; e->w = newW; break; }
for (Edge *e = g[v]; e; e = e->next)
if (e->to == u) { e->w = newW; break; }
int delta = newW - oldW;
int child = (tin[u] < tin[v]) ? v : u; /* tin 大的是子节点 */
segAdd(seg, 1, 1, n, tin[child], tout[child], delta);
} else { /* 查询到根距离 */
int x = queries[i][1];
ans[cnt++] = (int)segQuery(seg, 1, 1, n, tin[x]);
}
}
/* 释放内存 */
segFree(seg);
free(tin); free(tout); free(dist);
for (int i = 1; i <= n; i++) {
Edge *e = g[i];
while (e) { Edge *tmp = e; e = e->next; free(tmp); }
}
free(g);
*returnSize = cnt;
return ans;
}
```
---
关键点说明
要点 说明
迭代 DFS 显式栈实现,避免链状树递归深度 `1e5` 导致栈溢出。
判断父子 `tin[u] < tin[v]` 说明 `u` 先被访问,`v` 在 `u` 的子树中,因此 `v` 是子节点。
懒标记线段树 懒标记不下传,查询时累加根到叶子路径上的所有标记,天然支持区间加、单点查。
边权更新 直接遍历 `u`、`v` 的邻接表找到对方,读取旧值、更新新值,无需额外哈希数组。
---
复杂度
- 时间:`O((n + q) log n)` — DFS `O(n)`,每次查询/修改线段树 `O(log n)`。
- 空间:`O(n)` — 邻接表、线段树、时间戳数组等。

更多推荐

所有评论(0)