两种方法：懒删除堆 / 倒序处理（Python/Java/C++/Go）

LeetCode 每日一题题解

作者 endlesscheng

2025年7月6日 12:08

方法一：懒删除堆

首先，建图 + DFS，把每个连通块中的节点加到各自的最小堆中。每个最小堆维护对应连通块的节点编号。

然后处理询问。

对于类型二，用一个 $\textit{offline}$ 布尔数组表示离线的电站。这一步不修改堆。

对于类型一：

如果电站 $x$ 在线，那么答案为 $x$。
否则检查 $x$ 所处堆的堆顶是否在线。若离线，则弹出堆顶，重复该过程。如果堆为不空，那么答案为堆顶，否则为 $-1$。

为了找到 $x$ 所属的堆，还需要一个数组 $\textit{belong}$ 记录每个节点在哪个堆中。

具体请看视频讲解，欢迎点赞关注~

###py

class Solution:
    def processQueries(self, c: int, connections: List[List[int]], queries: List[List[int]]) -> List[int]:
        g = [[] for _ in range(c + 1)]
        for x, y in connections:
            g[x].append(y)
            g[y].append(x)

        belong = [-1] * (c + 1)
        heaps = []

        def dfs(x: int) -> None:
            belong[x] = len(heaps)  # 记录节点 x 在哪个堆
            h.append(x)
            for y in g[x]:
                if belong[y] < 0:
                    dfs(y)

        for i in range(1, c + 1):
            if belong[i] >= 0:
                continue
            h = []
            dfs(i)
            heapify(h)
            heaps.append(h)

        ans = []
        offline = [False] * (c + 1)
        for op, x in queries:
            if op == 2:
                offline[x] = True
                continue
            if not offline[x]:
                ans.append(x)
                continue
            h = heaps[belong[x]]
            # 懒删除：取堆顶的时候，如果离线，才删除
            while h and offline[h[0]]:
                heappop(h)
            ans.append(h[0] if h else -1)
        return ans

###java

class Solution {
    public int[] processQueries(int c, int[][] connections, int[][] queries) {
        List<Integer>[] g = new ArrayList[c + 1];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int[] e : connections) {
            int x = e[0], y = e[1];
            g[x].add(y);
            g[y].add(x);
        }

        int[] belong = new int[c + 1];
        Arrays.fill(belong, -1);
        List<PriorityQueue<Integer>> heaps = new ArrayList<>();
        PriorityQueue<Integer> pq;
        for (int i = 1; i <= c; i++) {
            if (belong[i] >= 0) {
                continue;
            }
            pq = new PriorityQueue<>();
            dfs(i, g, belong, heaps.size(), pq);
            heaps.add(pq);
        }

        int ansSize = 0;
        for (int[] q : queries) {
            if (q[0] == 1) {
                ansSize++;
            }
        }

        int[] ans = new int[ansSize];
        int idx = 0;
        boolean[] offline = new boolean[c + 1];
        for (int[] q : queries) {
            int x = q[1];
            if (q[0] == 2) {
                offline[x] = true;
                continue;
            }
            if (!offline[x]) {
                ans[idx++] = x;
                continue;
            }
            pq = heaps.get(belong[x]);
            // 懒删除：取堆顶的时候，如果离线，才删除
            while (!pq.isEmpty() && offline[pq.peek()]) {
                pq.poll();
            }
            ans[idx++] = pq.isEmpty() ? -1 : pq.peek();
        }
        return ans;
    }

    private void dfs(int x, List<Integer>[] g, int[] belong, int compId, PriorityQueue<Integer> pq) {
        belong[x] = compId; // 记录节点 x 在哪个堆
        pq.offer(x);
        for (int y : g[x]) {
            if (belong[y] < 0) {
                dfs(y, g, belong, compId, pq);
            }
        }
    }
}

###cpp

class Solution {
public:
    vector<int> processQueries(int c, vector<vector<int>>& connections, vector<vector<int>>& queries) {
        vector<vector<int>> g(c + 1);
        for (auto& e : connections) {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }

        vector<int> belong(c + 1, -1);
        vector<priority_queue<int, vector<int>, greater<>>> heaps;
        priority_queue<int, vector<int>, greater<>> pq;

        auto dfs = [&](this auto&& dfs, int x) -> void {
            belong[x] = heaps.size(); // 记录节点 x 在哪个堆
            pq.push(x);
            for (int y : g[x]) {
                if (belong[y] < 0) {
                    dfs(y);
                }
            }
        };

        for (int i = 1; i <= c; i++) {
            if (belong[i] < 0) {
                dfs(i);
                heaps.emplace_back(move(pq));
            }
        }

        vector<int> ans;
        vector<int8_t> offline(c + 1);
        for (auto& q : queries) {
            int x = q[1];
            if (q[0] == 2) {
                offline[x] = true;
                continue;
            }
            if (!offline[x]) {
                ans.push_back(x);
                continue;
            }
            auto& h = heaps[belong[x]];
            // 懒删除：取堆顶的时候，如果离线，才删除
            while (!h.empty() && offline[h.top()]) {
                h.pop();
            }
            ans.push_back(h.empty() ? -1 : h.top());
        }
        return ans;
    }
};

###go

func processQueries(c int, connections [][]int, queries [][]int) (ans []int) {
g := make([][]int, c+1)
for _, e := range connections {
x, y := e[0], e[1]
g[x] = append(g[x], y)
g[y] = append(g[y], x)
}

belong := make([]int, c+1)
for i := range belong {
belong[i] = -1
}
heaps := []hp{}
var h hp

var dfs func(int)
dfs = func(x int) {
belong[x] = len(heaps) // 记录节点 x 在哪个堆
h.IntSlice = append(h.IntSlice, x)
for _, y := range g[x] {
if belong[y] < 0 {
dfs(y)
}
}
}
for i := 1; i <= c; i++ {
if belong[i] >= 0 {
continue
}
h = hp{}
dfs(i)
heap.Init(&h)
heaps = append(heaps, h)
}

offline := make([]bool, c+1)
for _, q := range queries {
x := q[1]
if q[0] == 2 {
offline[x] = true
continue
}
if !offline[x] {
ans = append(ans, x)
continue
}
// 懒删除：取堆顶的时候，如果离线，才删除
h := &heaps[belong[x]]
for h.Len() > 0 && offline[h.IntSlice[0]] {
heap.Pop(h)
}
if h.Len() > 0 {
ans = append(ans, h.IntSlice[0])
} else {
ans = append(ans, -1)
}
}
return
}

type hp struct{ sort.IntSlice }
func (h *hp) Push(v any) { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() any   { a := h.IntSlice; v := a[len(a)-1]; h.IntSlice = a[:len(a)-1]; return v }

复杂度分析

时间复杂度：$\mathcal{O}(c\log c+n + q\log c)$ 或者 $\mathcal{O}(c+n + q\log c)$，取决于实现，其中 $n$ 是 $\textit{connections}$ 的长度，$q$ 是 $\textit{queries}$ 的长度。
空间复杂度：$\mathcal{O}(c+n)$。返回值不计入。

方法二：倒序处理 + 维护最小值

倒序处理询问，离线变成在线，删除变成添加，每个连通块只需要一个 $\texttt{int}$ 变量就可以维护最小值。

注意可能存在同一个节点多次离线的情况，我们需要记录节点离线的最早时间（询问的下标）。对于倒序处理来说，离线的最早时间才是真正的在线时间。

###py

class Solution:
    def processQueries(self, c: int, connections: List[List[int]], queries: List[List[int]]) -> List[int]:
        g = [[] for _ in range(c + 1)]
        for x, y in connections:
            g[x].append(y)
            g[y].append(x)

        belong = [-1] * (c + 1)
        cc = 0  # 连通块编号

        def dfs(x: int) -> None:
            belong[x] = cc  # 记录节点 x 在哪个连通块
            for y in g[x]:
                if belong[y] < 0:
                    dfs(y)

        for i in range(1, c + 1):
            if belong[i] < 0:
                dfs(i)
                cc += 1

        # 记录每个节点的离线时间，初始为无穷大（始终在线）
        offline_time = [inf] * (c + 1)
        for i in range(len(queries) - 1, -1, -1):
            t, x = queries[i]
            if t == 2:
                offline_time[x] = i  # 记录离线时间

        # 每个连通块中仍在线的电站的最小编号
        mn = [inf] * cc
        for i in range(1, c + 1):
            if offline_time[i] == inf:  # 最终仍在线
                j = belong[i]
                mn[j] = min(mn[j], i)

        ans = []
        for i in range(len(queries) - 1, -1, -1):
            t, x = queries[i]
            j = belong[x]
            if t == 2:
                if offline_time[x] == i:
                    mn[j] = min(mn[j], x)  # 变回在线
            elif i < offline_time[x]:  # 已经在线（写 < 或者 <= 都可以）
                ans.append(x)
            elif mn[j] != inf:
                ans.append(mn[j])
            else:
                ans.append(-1)
        ans.reverse()
        return ans

###java

class Solution {
    public int[] processQueries(int c, int[][] connections, int[][] queries) {
        List<Integer>[] g = new ArrayList[c + 1];
        Arrays.setAll(g, i -> new ArrayList<>());
        for (int[] e : connections) {
            int x = e[0], y = e[1];
            g[x].add(y);
            g[y].add(x);
        }

        int[] belong = new int[c + 1];
        Arrays.fill(belong, -1);
        int cc = 0; // 连通块编号
        for (int i = 1; i <= c; i++) {
            if (belong[i] < 0) {
                dfs(i, g, belong, cc);
                cc++;
            }
        }

        int[] offlineTime = new int[c + 1];
        Arrays.fill(offlineTime, Integer.MAX_VALUE);
        int q1 = 0;
        for (int i = queries.length - 1; i >= 0; i--) {
            int[] q = queries[i];
            if (q[0] == 2) {
                offlineTime[q[1]] = i; // 记录最早离线时间
            } else {
                q1++;
            }
        }

        // 维护每个连通块的在线电站的最小编号
        int[] mn = new int[cc];
        Arrays.fill(mn, Integer.MAX_VALUE);
        for (int i = 1; i <= c; i++) {
            if (offlineTime[i] == Integer.MAX_VALUE) { // 最终仍然在线
                int j = belong[i];
                mn[j] = Math.min(mn[j], i);
            }
        }

        int[] ans = new int[q1];
        for (int i = queries.length - 1; i >= 0; i--) {
            int[] q = queries[i];
            int x = q[1];
            int j = belong[x];
            if (q[0] == 2) {
                if (offlineTime[x] == i) { // 变回在线
                    mn[j] = Math.min(mn[j], x);
                }
            } else {
                q1--;
                if (i < offlineTime[x]) { // 已经在线（写 < 或者 <= 都可以）
                    ans[q1] = x;
                } else if (mn[j] != Integer.MAX_VALUE) {
                    ans[q1] = mn[j];
                } else {
                    ans[q1] = -1;
                }
            }
        }
        return ans;
    }

    private void dfs(int x, List<Integer>[] g, int[] belong, int compId) {
        belong[x] = compId;
        for (int y : g[x]) {
            if (belong[y] < 0) {
                dfs(y, g, belong, compId);
            }
        }
    }
}

###cpp

class Solution {
public:
    vector<int> processQueries(int c, vector<vector<int>>& connections, vector<vector<int>>& queries) {
        vector<vector<int>> g(c + 1);
        for (auto& e : connections) {
            int x = e[0], y = e[1];
            g[x].push_back(y);
            g[y].push_back(x);
        }

        vector<int> belong(c + 1, -1);
        int cc = 0; // 连通块编号
        auto dfs = [&](this auto&& dfs, int x) -> void {
            belong[x] = cc; // 记录节点 x 在哪个连通块
            for (int y : g[x]) {
                if (belong[y] < 0) {
                    dfs(y);
                }
            }
        };

        for (int i = 1; i <= c; i++) {
            if (belong[i] < 0) {
                dfs(i);
                cc++;
            }
        }

        vector<int> offline_time(c + 1, INT_MAX);
        for (int i = queries.size() - 1; i >= 0; i--) {
            auto& q = queries[i];
            if (q[0] == 2) {
                offline_time[q[1]] = i; // 记录最早离线时间
            }
        }

        // 维护每个连通块的在线电站的最小编号
        vector<int> mn(cc, INT_MAX);
        for (int i = 1; i <= c; i++) {
            if (offline_time[i] == INT_MAX) { // 最终仍然在线
                int j = belong[i];
                mn[j] = min(mn[j], i);
            }
        }

        vector<int> ans;
        for (int i = queries.size() - 1; i >= 0; i--) {
            auto& q = queries[i];
            int x = q[1];
            int j = belong[x];
            if (q[0] == 2) {
                if (offline_time[x] == i) { // 变回在线
                    mn[j] = min(mn[j], x);
                }
            } else if (i < offline_time[x]) { // 已经在线（写 < 或者 <= 都可以）
                ans.push_back(x);
            } else if (mn[j] != INT_MAX) {
                ans.push_back(mn[j]);
            } else {
                ans.push_back(-1);
            }
        }
        ranges::reverse(ans);
        return ans;
    }
};

###go

func processQueries(c int, connections [][]int, queries [][]int) []int {
g := make([][]int, c+1)
for _, e := range connections {
x, y := e[0], e[1]
g[x] = append(g[x], y)
g[y] = append(g[y], x)
}

belong := make([]int, c+1)
for i := range belong {
belong[i] = -1
}
cc := 0 // 连通块编号

var dfs func(int)
dfs = func(x int) {
belong[x] = cc // 记录节点 x 在哪个连通块
for _, y := range g[x] {
if belong[y] < 0 {
dfs(y)
}
}
}
for i := 1; i <= c; i++ {
if belong[i] < 0 {
dfs(i)
cc++
}
}

offlineTime := make([]int, c+1)
for i := range offlineTime {
offlineTime[i] = math.MaxInt
}
q1 := 0
for i, q := range slices.Backward(queries) {
if q[0] == 2 {
offlineTime[q[1]] = i // 记录最早离线时间
} else {
q1++
}
}

// 维护每个连通块的在线电站的最小编号
mn := make([]int, cc)
for i := range mn {
mn[i] = math.MaxInt
}
for i := 1; i <= c; i++ {
if offlineTime[i] == math.MaxInt { // 最终仍然在线
j := belong[i]
mn[j] = min(mn[j], i)
}
}

ans := make([]int, q1)
for i, q := range slices.Backward(queries) {
x := q[1]
j := belong[x]
if q[0] == 2 {
if offlineTime[x] == i { // 变回在线
mn[j] = min(mn[j], x)
}
} else {
q1--
if i < offlineTime[x] { // 已经在线（写 < 或者 <= 都可以）
ans[q1] = x
} else if mn[j] != math.MaxInt {
ans[q1] = mn[j]
} else {
ans[q1] = -1
}
}
}
return ans
}

复杂度分析

时间复杂度：$\mathcal{O}(c+n + q)$，其中 $n$ 是 $\textit{connections}$ 的长度，$q$ 是 $\textit{queries}$ 的长度。
空间复杂度：$\mathcal{O}(c+n)$。返回值不计入。

相似题目

3108. 带权图里旅途的最小代价

专题训练

图论题单的「§1.1 DFS 基础」。
数据结构题单的「§5.6 懒删除堆」。
数据结构题单的「专题：离线算法」。

分类题单

如何科学刷题？

我的题解精选（已分类）

LeetCode 每日一题题解
两个有序集合维护前 x 大二元组（Python/Java/C++/Go）endlesscheng
2024年10月13日 12:36

两个有序集合维护前 x 大二元组（Python/Java/C++/Go）

LeetCode 每日一题题解

作者 endlesscheng

2024年10月13日 12:36

前置题目：

在 3013 题中，我们用两个有序集合维护前 $k-1$ 小元素及其总和。

本题要维护前 $x$ 大的二元组 $(\textit{cnt}[x], x)$，以及 $\textit{cnt}[x]\cdot x$ 的总和。其中 $\textit{cnt}[x]$ 表示 $x$ 在子数组（滑动窗口）中的出现次数。

当元素进入窗口时：

把 $(\textit{cnt}[x], x)$ 从有序集合中移除。
把 $\textit{cnt}[x]$ 加一。
把 $(\textit{cnt}[x], x)$ 加入有序集合。

当元素离开窗口时：

把 $(\textit{cnt}[x], x)$ 从有序集合中移除。
把 $\textit{cnt}[x]$ 减一。
把 $(\textit{cnt}[x], x)$ 加入有序集合。

添加删除的同时维护 $\textit{cnt}[x]\cdot x$ 的总和。

其余逻辑同 3013 题。

本题视频讲解，欢迎点赞关注~

###py

from sortedcontainers import SortedList

class Solution:
    def findXSum(self, nums: List[int], k: int, x: int) -> List[int]:
        cnt = defaultdict(int)
        L = SortedList()  # 保存 tuple (出现次数，元素值)
        R = SortedList()
        sum_l = 0  # L 的元素和

        def add(val: int) -> None:
            if cnt[val] == 0:
                return
            p = (cnt[val], val)
            if L and p > L[0]:  # p 比 L 中最小的还大
                nonlocal sum_l
                sum_l += p[0] * p[1]
                L.add(p)
            else:
                R.add(p)

        def remove(val: int) -> None:
            if cnt[val] == 0:
                return
            p = (cnt[val], val)
            if p in L:
                nonlocal sum_l
                sum_l -= p[0] * p[1]
                L.remove(p)
            else:
                R.remove(p)

        def l2r() -> None:
            nonlocal sum_l
            p = L[0]
            sum_l -= p[0] * p[1]
            L.remove(p)
            R.add(p)

        def r2l() -> None:
            nonlocal sum_l
            p = R[-1]
            sum_l += p[0] * p[1]
            R.remove(p)
            L.add(p)

        ans = [0] * (len(nums) - k + 1)
        for r, in_ in enumerate(nums):
            # 添加 in_
            remove(in_)
            cnt[in_] += 1
            add(in_)

            l = r + 1 - k
            if l < 0:
                continue

            # 维护大小
            while R and len(L) < x:
                r2l()
            while len(L) > x:
                l2r()
            ans[l] = sum_l

            # 移除 out
            out = nums[l]
            remove(out)
            cnt[out] -= 1
            add(out)
        return ans

###java

class Solution {
    private final TreeSet<int[]> L = new TreeSet<>((a, b) -> a[0] != b[0] ? a[0] - b[0] : a[1] - b[1]);
    private final TreeSet<int[]> R = new TreeSet<>(L.comparator());
    private final Map<Integer, Integer> cnt = new HashMap<>();
    private long sumL = 0;

    public long[] findXSum(int[] nums, int k, int x) {
        long[] ans = new long[nums.length - k + 1];
        for (int r = 0; r < nums.length; r++) {
            // 添加 in
            int in = nums[r];
            del(in);
            cnt.merge(in, 1, Integer::sum); // cnt[in]++
            add(in);

            int l = r + 1 - k;
            if (l < 0) {
                continue;
            }

            // 维护大小
            while (!R.isEmpty() && L.size() < x) {
                r2l();
            }
            while (L.size() > x) {
                l2r();
            }
            ans[l] = sumL;

            // 移除 out
            int out = nums[l];
            del(out);
            cnt.merge(out, -1, Integer::sum); // cnt[out]--
            add(out);
        }
        return ans;
    }

    // 添加元素
    private void add(int val) {
        int c = cnt.get(val);
        if (c == 0) {
            return;
        }
        int[] p = new int[]{c, val};
        if (!L.isEmpty() && L.comparator().compare(p, L.first()) > 0) { // p 比 L 中最小的还大
            sumL += (long) p[0] * p[1];
            L.add(p);
        } else {
            R.add(p);
        }
    }

    // 删除元素
    private void del(int val) {
        int c = cnt.getOrDefault(val, 0);
        if (c == 0) {
            return;
        }
        int[] p = new int[]{c, val};
        if (L.contains(p)) {
            sumL -= (long) p[0] * p[1];
            L.remove(p);
        } else {
            R.remove(p);
        }
    }

    // 从 L 移动一个元素到 R
    private void l2r() {
        int[] p = L.pollFirst();
        sumL -= (long) p[0] * p[1];
        R.add(p);
    }

    // 从 R 移动一个元素到 L
    private void r2l() {
        int[] p = R.pollLast();
        sumL += (long) p[0] * p[1];
        L.add(p);
    }
}

###cpp

class Solution {
public:
    vector<long long> findXSum(vector<int>& nums, int k, int x) {
        using pii = pair<int, int>; // 出现次数，元素值
        set<pii> L, R;
        long long sum_l = 0; // L 的元素和
        unordered_map<int, int> cnt;
        auto add = [&](int x) {
            pii p = {cnt[x], x};
            if (p.first == 0) {
                return;
            }
            if (!L.empty() && p > *L.begin()) { // p 比 L 中最小的还大
                sum_l += (long long) p.first * p.second;
                L.insert(p);
            } else {
                R.insert(p);
            }
        };
        auto del = [&](int x) {
            pii p = {cnt[x], x};
            if (p.first == 0) {
                return;
            }
            auto it = L.find(p);
            if (it != L.end()) {
                sum_l -= (long long) p.first * p.second;
                L.erase(it);
            } else {
                R.erase(p);
            }
        };
        auto l2r = [&]() {
            pii p = *L.begin();
            sum_l -= (long long) p.first * p.second;
            L.erase(p);
            R.insert(p);
        };
        auto r2l = [&]() {
            pii p = *R.rbegin();
            sum_l += (long long) p.first * p.second;
            R.erase(p);
            L.insert(p);
        };

        vector<long long> ans(nums.size() - k + 1);
        for (int r = 0; r < nums.size(); r++) {
            // 添加 in
            int in = nums[r];
            del(in);
            cnt[in]++;
            add(in);

            int l = r + 1 - k;
            if (l < 0) {
                continue;
            }

            // 维护大小
            while (!R.empty() && L.size() < x) {
                r2l();
            }
            while (L.size() > x) {
                l2r();
            }
            ans[l] = sum_l;

            // 移除 out
            int out = nums[l];
            del(out);
            cnt[out]--;
            add(out);
        }
        return ans;
    }
};

###go

import "github.com/emirpasic/gods/v2/trees/redblacktree"

type pair struct{ c, x int } // 出现次数，元素值

func less(p, q pair) int {
return cmp.Or(p.c-q.c, p.x-q.x)
}

func findXSum(nums []int, k, x int) []int64 {
L := redblacktree.NewWith[pair, struct{}](less)
R := redblacktree.NewWith[pair, struct{}](less)

sumL := 0 // L 的元素和
cnt := map[int]int{}
add := func(x int) {
p := pair{cnt[x], x}
if p.c == 0 {
return
}
if !L.Empty() && less(p, L.Left().Key) > 0 { // p 比 L 中最小的还大
sumL += p.c * p.x
L.Put(p, struct{}{})
} else {
R.Put(p, struct{}{})
}
}
del := func(x int) {
p := pair{cnt[x], x}
if p.c == 0 {
return
}
if _, ok := L.Get(p); ok {
sumL -= p.c * p.x
L.Remove(p)
} else {
R.Remove(p)
}
}
l2r := func() {
p := L.Left().Key
sumL -= p.c * p.x
L.Remove(p)
R.Put(p, struct{}{})
}
r2l := func() {
p := R.Right().Key
sumL += p.c * p.x
R.Remove(p)
L.Put(p, struct{}{})
}

ans := make([]int64, len(nums)-k+1)
for r, in := range nums {
// 添加 in
del(in)
cnt[in]++
add(in)

l := r + 1 - k
if l < 0 {
continue
}

// 维护大小
for !R.Empty() && L.Size() < x {
r2l()
}
for L.Size() > x {
l2r()
}
ans[l] = int64(sumL)

// 移除 out
out := nums[l]
del(out)
cnt[out]--
add(out)
}
return ans
}

复杂度分析

时间复杂度：$\mathcal{O}(n\log k)$，其中 $n$ 是 $\textit{nums}$ 的长度。
空间复杂度：$\mathcal{O}(n)$。

专题训练

见下面数据结构题单的「§5.7 对顶堆」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
贪心，每段保留最大的（Python/Java/C++/C/Go/JS/Rust）endlesscheng
2025年10月17日 17:06

贪心，每段保留最大的（Python/Java/C++/C/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2025年10月17日 17:06

为了不让相邻气球颜色相同，对于 $\textit{colors}$ 的每个连续同色段，只能保留一个气球。

贪心地，保留其中耗时最大的气球。

答案为 $\textit{neededTime}$ 的总和，减去每段的最大耗时。

###py

class Solution:
    def minCost(self, colors: str, neededTime: List[int]) -> int:
        ans = max_t = 0
        for i, t in enumerate(neededTime):
            ans += t
            if t > max_t:  # 手写 if 比调用 max 快
                max_t = t
            if i == len(colors) - 1 or colors[i] != colors[i + 1]:
                # 遍历到了连续同色段的末尾
                ans -= max_t  # 保留耗时最大的气球
                max_t = 0  # 准备计算下一段的最大耗时
        return ans

###java

class Solution {
    public int minCost(String colors, int[] neededTime) {
        int n = neededTime.length;
        int ans = 0;
        int maxT = 0;
        for (int i = 0; i < n; i++) {
            int t = neededTime[i];
            ans += t;
            maxT = Math.max(maxT, t);
            if (i == n - 1 || colors.charAt(i) != colors.charAt(i + 1)) {
                // 遍历到了连续同色段的末尾
                ans -= maxT; // 保留耗时最大的气球
                maxT = 0; // 准备计算下一段的最大耗时
            }
        }
        return ans;
    }
}

###cpp

class Solution {
public:
    int minCost(string colors, vector<int>& neededTime) {
        int n = colors.size();
        int ans = 0, max_t = 0;
        for (int i = 0; i < n; i++) {
            int t = neededTime[i];
            ans += t;
            max_t = max(max_t, t);
            if (i == n - 1 || colors[i] != colors[i + 1]) {
                // 遍历到了连续同色段的末尾
                ans -= max_t; // 保留耗时最大的气球
                max_t = 0; // 准备计算下一段的最大耗时
            }
        }
        return ans;
    }
};

###c

#define MAX(a, b) ((b) > (a) ? (b) : (a))

int minCost(char* colors, int* neededTime, int neededTimeSize) {
    int ans = 0, max_t = 0;
    for (int i = 0; i < neededTimeSize; i++) {
        int t = neededTime[i];
        ans += t;
        max_t = MAX(max_t, t);
        if (i == neededTimeSize - 1 || colors[i] != colors[i + 1]) {
            // 遍历到了连续同色段的末尾
            ans -= max_t; // 保留耗时最大的气球
            max_t = 0; // 准备计算下一段的最大耗时
        }
    }
    return ans;
}

###go

func minCost(colors string, neededTime []int) (ans int) {
maxT := 0
for i, t := range neededTime {
ans += t
maxT = max(maxT, t)
if i == len(colors)-1 || colors[i] != colors[i+1] {
// 遍历到了连续同色段的末尾
ans -= maxT // 保留耗时最大的气球
maxT = 0    // 准备计算下一段的最大耗时
}
}
return
}

###js

var minCost = function(colors, neededTime) {
    const n = colors.length;
    let ans = 0, maxT = 0;
    for (let i = 0; i < n; i++) {
        const t = neededTime[i];
        ans += t;
        maxT = Math.max(maxT, t);
        if (i === n - 1 || colors[i] !== colors[i + 1]) {
            // 遍历到了连续同色段的末尾
            ans -= maxT; // 保留耗时最大的气球
            maxT = 0; // 准备计算下一段的最大耗时
        }
    }
    return ans;
};

###rust

impl Solution {
    pub fn min_cost(colors: String, needed_time: Vec<i32>) -> i32 {
        let s = colors.as_bytes();
        let mut ans = 0;
        let mut max_t = 0;
        for (i, t) in needed_time.into_iter().enumerate() {
            ans += t;
            max_t = max_t.max(t);
            if i + 1 == s.len() || s[i] != s[i + 1] {
                // 遍历到了连续同色段的末尾
                ans -= max_t; // 保留耗时最大的气球
                max_t = 0; // 准备计算下一段的最大耗时
            }
        }
        ans
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(n)$，其中 $n$ 是 $\textit{colors}$ 的长度。
空间复杂度：$\mathcal{O}(1)$。

专题训练

见下面双指针题单的「六、分组循环」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
哨兵节点+一次遍历（Python/Java/C++/C/Go/JS/Rust）endlesscheng
2024年7月14日 12:10

哨兵节点+一次遍历（Python/Java/C++/C/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2024年7月14日 12:10

如何在遍历链表的同时，删除链表节点？请看【基础算法精讲 08】。

对于本题，由于直接判断节点值是否在 $\textit{nums}$ 中，需要遍历 $\textit{nums}$，时间复杂度为 $\mathcal{O}(n)$。把 $\textit{nums}$ 中的元素保存一个哈希集合中，然后判断节点值是否在哈希集合中，这样可以做到 $\mathcal{O}(1)$。

具体做法：

把 $\textit{nums}$ 中的元素保存到一个哈希集合中。
由于头节点可能会被删除，在头节点前面插入一个哨兵节点 $\textit{dummy}$，以简化代码逻辑。
初始化 $\textit{cur} = \textit{dummy}$。
遍历链表，如果 $\textit{cur}$ 的下一个节点的值在哈希集合中，则需要删除，更新 $\textit{cur}.\textit{next}$ 为 $\textit{cur}.\textit{next}.\textit{next}$；否则不删除，更新 $\textit{cur}$ 为 $\textit{cur}.\textit{next}$。
循环结束后，返回 $\textit{dummy}.\textit{next}$。

注：$\textit{dummy}$ 和 $\textit{cur}$ 是同一个节点的引用，修改 $\textit{cur}.\textit{next}$ 也会修改 $\textit{dummy}.\textit{next}$。

本题视频讲解，欢迎点赞关注~

###py

class Solution:
    def modifiedList(self, nums: List[int], head: Optional[ListNode]) -> Optional[ListNode]:
        st = set(nums)
        cur = dummy = ListNode(next=head)
        while cur.next:
            nxt = cur.next
            if nxt.val in st:
                cur.next = nxt.next  # 从链表中删除 nxt 节点
            else:
                cur = nxt  # 不删除 nxt，继续向后遍历链表
        return dummy.next

###java

class Solution {
    public ListNode modifiedList(int[] nums, ListNode head) {
        Set<Integer> set = new HashSet<>(nums.length, 1); // 预分配空间
        for (int x : nums) {
            set.add(x);
        }

        ListNode dummy = new ListNode(0, head);
        ListNode cur = dummy;
        while (cur.next != null) {
            ListNode nxt = cur.next;
            if (set.contains(nxt.val)) {
                cur.next = nxt.next; // 从链表中删除 nxt 节点
            } else {
                cur = nxt; // 不删除 nxt，继续向后遍历链表
            }
        }
        return dummy.next;
    }
}

###cpp

class Solution {
public:
    ListNode* modifiedList(vector<int>& nums, ListNode* head) {
        unordered_set<int> st(nums.begin(), nums.end());
        ListNode dummy(0, head);
        ListNode* cur = &dummy;
        while (cur->next) {
            ListNode* nxt = cur->next;
            if (st.contains(nxt->val)) {
                cur->next = nxt->next; // 从链表中删除 nxt 节点
            } else {
                cur = nxt; // 不删除 nxt，继续向后遍历链表
            }
        }
        return dummy.next;
    }
};

###c

#define MAX(a, b) ((b) > (a) ? (b) : (a))

struct ListNode* modifiedList(int* nums, int numsSize, struct ListNode* head) {
    int mx = 0;
    for (int i = 0; i < numsSize; i++) {
        mx = MAX(mx, nums[i]);
    }

    bool* has = calloc(mx + 1, sizeof(bool));
    for (int i = 0; i < numsSize; i++) {
        has[nums[i]] = true;
    }

    struct ListNode dummy = {0, head};
    struct ListNode* cur = &dummy;
    while (cur->next) {
        struct ListNode* nxt = cur->next;
        if (nxt->val <= mx && has[nxt->val]) {
            cur->next = nxt->next; // 从链表中删除 nxt 节点
            free(nxt);
        } else {
            cur = nxt; // 不删除 nxt，继续向后遍历链表
        }
    }

    free(has);
    return dummy.next;
}

###go

func modifiedList(nums []int, head *ListNode) *ListNode {
has := make(map[int]bool, len(nums)) // 预分配空间
for _, x := range nums {
has[x] = true
}

dummy := &ListNode{Next: head}
cur := dummy
for cur.Next != nil {
nxt := cur.Next
if has[nxt.Val] {
cur.Next = nxt.Next // 从链表中删除 nxt 节点
} else {
cur = nxt // 不删除 nxt，继续向后遍历链表
}
}
return dummy.Next
}

###js

var modifiedList = function(nums, head) {
    const set = new Set(nums);
    const dummy = new ListNode(0, head);
    let cur = dummy;
    while (cur.next) {
        const nxt = cur.next;
        if (set.has(nxt.val)) {
            cur.next = nxt.next; // 从链表中删除 nxt 节点
        } else {
            cur = nxt; // 不删除 nxt，继续向后遍历链表
        }
    }
    return dummy.next;
};

###rust

use std::collections::HashSet;

impl Solution {
    pub fn modified_list(nums: Vec<i32>, head: Option<Box<ListNode>>) -> Option<Box<ListNode>> {
        let set = nums.into_iter().collect::<HashSet<_>>();
        let mut dummy = Box::new(ListNode { val: 0, next: head });
        let mut cur = &mut dummy;
        while let Some(ref mut nxt) = cur.next {
            if set.contains(&nxt.val) {
                cur.next = nxt.next.take(); // 从链表中删除 nxt 节点
            } else {
                cur = cur.next.as_mut()?; // 不删除 nxt，继续向后遍历链表
            }
        }
        dummy.next
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(n + m)$，其中 $n$ 是 $\textit{nums}$ 的长度，$m$ 是链表的长度。
空间复杂度：$\mathcal{O}(n)$。

专题训练

见下面链表题单的「§1.2 删除节点」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
O(1) 空间做法：位运算 / 数学（Python/Java/C++/Go）endlesscheng
2024年9月15日 12:09

O(1) 空间做法：位运算 / 数学（Python/Java/C++/Go）

LeetCode 每日一题题解

作者 endlesscheng

2024年9月15日 12:09

本题和 2965. 找出缺失和重复的数字本质是一样的，见我的题解，有位运算和数学两种做法。

位运算

需要两次遍历。一次遍历见下面的数学做法。

###py

class Solution:
    def getSneakyNumbers(self, nums: List[int]) -> List[int]:
        n = len(nums) - 2
        xor_all = n ^ (n + 1)  # n 和 n+1 多异或了
        for i, x in enumerate(nums):
            xor_all ^= i ^ x
        shift = xor_all.bit_length() - 1

        ans = [0, 0]
        for i, x in enumerate(nums):
            if i < n:
                ans[i >> shift & 1] ^= i
            ans[x >> shift & 1] ^= x
        return ans

###java

class Solution {
    public int[] getSneakyNumbers(int[] nums) {
        int n = nums.length - 2;
        int xorAll = n ^ (n + 1); // n 和 n+1 多异或了
        for (int i = 0; i < nums.length; i++) {
            xorAll ^= i ^ nums[i];
        }
        int shift = Integer.numberOfTrailingZeros(xorAll);

        int[] ans = new int[2];
        for (int i = 0; i < nums.length; i++) {
            if (i < n) {
                ans[i >> shift & 1] ^= i;
            }
            ans[nums[i] >> shift & 1] ^= nums[i];
        }
        return ans;
    }
}

###cpp

class Solution {
public:
    vector<int> getSneakyNumbers(vector<int>& nums) {
        int n = nums.size() - 2;
        int xor_all = n ^ (n + 1); // n 和 n+1 多异或了
        for (int i = 0; i < nums.size(); i++) {
            xor_all ^= i ^ nums[i];
        }
        int shift = __builtin_ctz(xor_all);

        vector<int> ans(2);
        for (int i = 0; i < nums.size(); i++) {
            if (i < n) {
                ans[i >> shift & 1] ^= i;
            }
            ans[nums[i] >> shift & 1] ^= nums[i];
        }
        return ans;
    }
};

###go

func getSneakyNumbers(nums []int) []int {
n := len(nums) - 2
xorAll := n ^ (n + 1) // n 和 n+1 多异或了
for i, x := range nums {
xorAll ^= i ^ x
}
shift := bits.TrailingZeros(uint(xorAll))

ans := make([]int, 2)
for i, x := range nums {
if i < n {
ans[i>>shift&1] ^= i
}
ans[x>>shift&1] ^= x
}
return ans
}

复杂度分析

时间复杂度：$\mathcal{O}(n)$，其中 $n$ 是 $\textit{nums}$ 的长度。
空间复杂度：$\mathcal{O}(1)$。

数学

设多出的两个数分别为 $x$ 和 $y$。

也就是说，$\textit{nums} = [0,1,2,\cdots,n-1,x,y]$。

设 $\textit{nums}$ 的元素和为 $s$，$\textit{nums}$ 的元素平方之和为 $s_2$，那么有

$$
\begin{aligned}
&x+y = s - (0 + 1 + 2 + \cdots + n-1) = a \
&x^2+y^2 = s_2 - (0^2 + 1^2 + 2^2 + \cdots + (n-1)^2) = b \
\end{aligned}
$$

解得

$$
\begin{cases}
x = \dfrac{a-\sqrt{2b-a^2}}{2} \
y = \dfrac{a+\sqrt{2b-a^2}}{2} \
\end{cases}
$$

也可以先算出 $x$，然后算出 $y=a-x$。

本题视频讲解，欢迎点赞关注~

###py

class Solution:
    def getSneakyNumbers(self, nums: List[int]) -> List[int]:
        n = len(nums) - 2
        a = -n * (n - 1) // 2
        b = -n * (n - 1) * (n * 2 - 1) // 6
        for x in nums:
            a += x
            b += x * x
        x = (a - isqrt(b * 2 - a * a)) // 2
        return [x, a - x]

###java

class Solution {
    public int[] getSneakyNumbers(int[] nums) {
        int n = nums.length - 2;
        int a = -n * (n - 1) / 2;
        int b = -n * (n - 1) * (n * 2 - 1) / 6;
        for (int x : nums) {
            a += x;
            b += x * x;
        }
        int x = (a - (int) Math.sqrt(b * 2 - a * a)) / 2;
        return new int[]{x, a - x};
    }
}

###cpp

class Solution {
public:
    vector<int> getSneakyNumbers(vector<int>& nums) {
        int n = nums.size() - 2;
        int a = -n * (n - 1) / 2;
        int b = -n * (n - 1) * (n * 2 - 1) / 6;
        for (int x : nums) {
            a += x;
            b += x * x;
        }
        int x = (a - (int) sqrt(b * 2 - a * a)) / 2;
        return {x, a - x};
    }
};

###go

func getSneakyNumbers(nums []int) []int {
n := len(nums) - 2
a := -n * (n - 1) / 2
b := -n * (n - 1) * (n*2 - 1) / 6
for _, x := range nums {
a += x
b += x * x
}
x := (a - int(math.Sqrt(float64(b*2-a*a)))) / 2
return []int{x, a - x}
}

复杂度分析

时间复杂度：$\mathcal{O}(n)$，其中 $n$ 是 $\textit{nums}$ 的长度。
空间复杂度：$\mathcal{O}(1)$。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
阅读理解（Python/Java/C++/C/Go/JS/Rust）endlesscheng
2022年1月2日 12:13

阅读理解（Python/Java/C++/C/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2022年1月2日 12:13

题意：去掉没有安全设备的行，计算相邻行之间的激光束的数量之和。

{:width=250px}

示例 1 有 $4$ 行，安全设备的个数分别为 $3,0,2,1$。

去掉没有安全设备的行，剩下 $3,2,1$。

计算相邻行之间的激光束的数量之和，即 $3\times 2+ 2\times 1 = 8$。

class Solution:
    def numberOfBeams(self, bank: List[str]) -> int:
        ans = pre_cnt = 0
        for row in bank:
            cnt = row.count('1')
            if cnt > 0:
                ans += pre_cnt * cnt
                pre_cnt = cnt
        return ans

class Solution {
    public int numberOfBeams(String[] bank) {
        int ans = 0;
        int preCnt = 0;
        for (String row : bank) {
            int cnt = 0;
            for (char ch : row.toCharArray()) {
                cnt += ch - '0';
            }
            if (cnt > 0) {
                ans += preCnt * cnt;
                preCnt = cnt;
            }
        }
        return ans;
    }
}

class Solution {
public:
    int numberOfBeams(vector<string>& bank) {
        int ans = 0, pre_cnt = 0;
        for (auto& row : bank) {
            int cnt = ranges::count(row, '1');
            if (cnt > 0) {
                ans += pre_cnt * cnt;
                pre_cnt = cnt;
            }
        }
        return ans;
    }
};

int numberOfBeams(char **bank, int bankSize) {
    int ans = 0, pre_cnt = 0;
    for (int i = 0; i < bankSize; i++) {
        char* row = bank[i];
        int cnt = 0;
        for (int j = 0; row[j]; j++) {
            cnt += row[j] - '0';
        }
        if (cnt > 0) {
            ans += pre_cnt * cnt;
            pre_cnt = cnt;
        }
    }
    return ans;
}

func numberOfBeams(bank []string) (ans int) {
preCnt := 0
for _, row := range bank {
cnt := strings.Count(row, "1")
if cnt > 0 {
ans += preCnt * cnt
preCnt = cnt
}
}
return
}

var numberOfBeams = function(bank) {
    let ans = 0, preCnt = 0;
    for (const row of bank) {
        const cnt = row.split('1').length - 1;
        if (cnt > 0) {
            ans += preCnt * cnt;
            preCnt = cnt;
        }
    }
    return ans;
};

impl Solution {
    pub fn number_of_beams(bank: Vec<String>) -> i32 {
        let mut ans = 0;
        let mut pre_cnt = 0;
        for row in bank {
            let cnt = row.bytes().filter(|&c| c == b'1').count() as i32;
            if cnt > 0 {
                ans += pre_cnt * cnt;
                pre_cnt = cnt;
            }
        }
        ans
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(mn)$，其中 $m$ 和 $n$ 分别是 $\textit{bank}$ 的行数和列数。
空间复杂度：$\mathcal{O}(1)$。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
O(1) 等差数列求和（Python/Java/C++/C/Go/JS/Rust）endlesscheng
2025年10月25日 08:09

O(1) 等差数列求和（Python/Java/C++/C/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2025年10月25日 08:09

设一周的天数为 $D=7$。

根据题意，第一周存的钱数为

$$
\textit{base} = 1+2+\cdots+D = \dfrac{D(D+1)}{2}
$$

第二周每天存的钱都比 $D$ 天前多 $1$，所以第二周存的钱数为

$$
\textit{base} + D
$$

第三周每天存的钱都比 $D$ 天前多 $1$，所以第三周存的钱数为

$$
\textit{base} + 2D
$$

依此类推。我们存了完整的 $w=\left\lfloor\dfrac{n}{D}\right\rfloor$ 周，在第 $w$ 周存的钱数为

$$
\textit{base} + (w-1)\cdot D
$$

$w$ 周一共存的钱数为

$$
\begin{aligned}
\sum_{i=0}^{w-1} \textit{base} + i\cdot D &= w\cdot \textit{base} + \dfrac{w(w-1)}{2}\cdot D \
&= w\cdot \dfrac{D(D+1)}{2} + \dfrac{w(w-1)}{2}\cdot D \
&= \dfrac{wD(w+D)}{2} \
\end{aligned}
$$

如果 $n$ 不是 $D$ 的倍数，还有 $r = n\bmod D$ 天，存的钱数为

$$
(w+1) + (w+2) + \cdots + (w+r) = rw + \dfrac{r(r+1)}{2} = \dfrac{r(2w+r+1)}{2}
$$

最终答案为

$$
\dfrac{wD(w+D)}{2} + \dfrac{r(2w+r+1)}{2} = \dfrac{wD(w+D) + r(2w+r+1)}{2}
$$

其中 $D=7$，$w=\left\lfloor\dfrac{n}{D}\right\rfloor$，$r = n\bmod D$。

class Solution:
    def totalMoney(self, n: int) -> int:
        D = 7
        w, r = divmod(n, D)
        return (w * D * (w + D) + r * (w * 2 + r + 1)) // 2

class Solution {
    public int totalMoney(int n) {
        final int D = 7;
        int w = n / D;
        int r = n % D;
        return (w * D * (w + D) + r * (w * 2 + r + 1)) / 2;
    }
}

class Solution {
public:
    int totalMoney(int n) {
        constexpr int D = 7;
        int w = n / D, r = n % D;
        return (w * D * (w + D) + r * (w * 2 + r + 1)) / 2;
    }
};

int totalMoney(int n) {
    const int D = 7;
    int w = n / D, r = n % D;
    return (w * D * (w + D) + r * (w * 2 + r + 1)) / 2;
}

func totalMoney(n int) int {
const d = 7
w, r := n/d, n%d
return (w*d*(w+d) + r*(w*2+r+1)) / 2
}

var totalMoney = function(n) {
    const D = 7;
    const w = Math.floor(n / D), r = n % D;
    return (w * D * (w + D) + r * (w * 2 + r + 1)) / 2;
};

impl Solution {
    pub fn total_money(n: i32) -> i32 {
        const D: i32 = 7;
        let w = n / D;
        let r = n % D;
        (w * D * (w + D) + r * (w * 2 + r + 1)) / 2
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(1)$。
空间复杂度：$\mathcal{O}(1)$。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
两种方法：枚举 / 倒序贪心 + 0-1 背包（Python/Java/C++/Go）endlesscheng
2025年10月24日 10:36

两种方法：枚举 / 倒序贪心 + 0-1 背包（Python/Java/C++/Go）

LeetCode 每日一题题解

作者 endlesscheng

2025年10月24日 10:36

方法一：枚举

设 $n$ 的十进制长度为 $m$。

对于本题的数据范围，一定存在十进制长为 $m+1$ 的数值平衡数。

例如 $n=999$，答案为 $1333$。

例如 $n=999999$，答案为 $1224444$。

对于本题的数据范围，$m+1$ 一定可以分解为 $[1,9]$ 中的不同元素之和。

所以枚举 $\mathcal{O}(n)$ 次就能找到答案。

###py

class Solution:
    def nextBeautifulNumber(self, n: int) -> int:
        while True:
            n += 1
            cnt = Counter(str(n))
            if all(int(d) == c for d, c in cnt.items()):
                return n

###java

class Solution {
    public int nextBeautifulNumber(int n) {
        next:
        while (true) {
            n++;

            int[] cnt = new int[10];
            for (int x = n; x > 0; x /= 10) {
                cnt[x % 10]++;
            }

            for (int x = n; x > 0; x /= 10) {
                if (cnt[x % 10] != x % 10) {
                    continue next;
                }
            }

            return n;
        }
    }
}

###cpp

class Solution {
public:
    int nextBeautifulNumber(int n) {
        while (true) {
            n++;

            int cnt[10]{};
            for (int x = n; x > 0; x /= 10) {
                cnt[x % 10]++;
            }

            bool ok = true;
            for (int x = n; x > 0; x /= 10) {
                if (cnt[x % 10] != x % 10) {
                    ok = false;
                    break;
                }
            }
            if (ok) {
                return n;
            }
        }
    }
};

###go

func nextBeautifulNumber(n int) int {
next:
for {
n++
cnt := [10]int{}
for x := n; x > 0; x /= 10 {
cnt[x%10]++
}
for x := n; x > 0; x /= 10 {
if cnt[x%10] != x%10 {
continue next
}
}
return n
}
}

复杂度分析

时间复杂度：$\mathcal{O}(n\log n)$ 或 $\mathcal{O}(n(D + \log n))$，其中 $D=10$。枚举 $\mathcal{O}(n)$ 个数，每个数需要 $\mathcal{O}(\log n)$ 的时间判断是否为数值平衡数。
空间复杂度：$\mathcal{O}(\log n)$ 或 $\mathcal{O}(D)$。

方法二：倒序贪心

如果数据范围扩大到 $n\le 10^{18}$，上面的做法就超时了。

下面介绍一个更快的做法，对于更大的数据范围也能瞬间得出结果，且可以扩展到其他情况，例如值域改成小写字母 $\texttt{a}$ 到 $\texttt{z}$。

请先完成更简单的 3720. 大于目标字符串的最小字典序排列，我的题解。

本题用同样的方法解决。

先把 $n$ 转成十进制字符串 $s$。为简化代码，在 $s$ 前面加一个前导零。设 $s$（补前导零后）的长度为 $m$。

倒着遍历 $s$，设 $d = s[i]$，尝试把 $s[i]$ 增大为 $d+1,d+2,\ldots,9$。那么：

对于下标在 $[1,i]$ 中的数字，不能存在数字 $x$ 的出现次数超过 $x$ 的情况。换句话说，设 $\textit{cnt}[x]$ 是数字 $x$ 在 $s$ 的 $[1,i]$ 中的出现次数，不能出现 $\textit{cnt}[x] > x$ 的情况。
剩余位置 $[i+1,m-1]$ 可以随便填，我们需要补满 $0 < \textit{cnt}[x] < x$ 的数字 $x$，把 $x$ 的出现次数变成 $x$。
补满数字后，如果还有剩余位置没有填数字，那么从剩余的满足 $\textit{cnt}[x] = 0$ 的非零数字中，选择一个字典序最小的序列，使得序列的和恰好等于剩余位置的个数。这是恰好装满型 0-1 背包。算完 DP 后，如何得到具体方案？类似 1449. 数位成本和为目标值的最大数字，见我的题解。
把第二步和第三步的数字从小到大排序，填在 $[i+1,m-1]$ 中。

###py

class Solution:
    # 从 a 中选一个字典序最小的、元素和等于 target 的子序列
    # a 已经从小到大排序
    # 无解返回 None
    def zeroOneKnapsack(self, a: List[int], target: int) -> Optional[List[int]]:
        n = len(a)
        f = [[False] * (target + 1) for _ in range(n + 1)]
        f[n][0] = True

        # 倒着 DP，这样后面可以正着（从小到大）选
        for i in range(n - 1, -1, -1):
            v = a[i]
            for j in range(target + 1):
                if j < v:
                    f[i][j] = f[i + 1][j]
                else:
                    f[i][j] = f[i + 1][j] or f[i + 1][j - v]

        if not f[0][target]:
            return None

        ans = []
        j = target
        for i, v in enumerate(a):
            if j >= v and f[i + 1][j - v]:
                ans.append(v)
                j -= v
        return ans

    def nextBeautifulNumber(self, n: int) -> int:
        s = "0" + str(n)  # 补一个前导零，方便处理答案十进制串比 n 的十进制串长的情况
        s = list(map(int, s))  # 避免在后续循环中反复调用 int
        m = len(s)

        MX = 10
        cnt = [0] * MX
        for i in range(1, m):
            cnt[s[i]] += 1

        # 从右往左尝试
        for i in range(m - 1, -1, -1):
            if i > 0:
                cnt[s[i]] -= 1  # 撤销

            # 增大 s[i] 为 j
            for j in range(s[i] + 1, MX):
                cnt[j] += 1

                # 后面 [i+1, m-1] 需要补满 0 < cnt[k] < k 的数字 k，剩余数位可以随便填
                free = m - 1 - i  # 统计可以随便填的数位个数
                for k, c in enumerate(cnt):
                    if k < c:  # 不合法
                        free = -1
                        break
                    if c > 0:
                        free -= k - c
                if free < 0:  # 不合法，继续枚举
                    cnt[j] -= 1
                    continue

                # 对于可以随便填的数位，计算字典序最小的填法
                a = [k for k in range(1, MX) if cnt[k] == 0]
                missing = self.zeroOneKnapsack(a, free)
                if missing is None:  # 无解，继续枚举
                    cnt[j] -= 1
                    continue

                for v in missing:
                    cnt[v] = -v  # 用负数表示可以随便填的数

                s[i] = j
                del s[i + 1:]
                for k, c in enumerate(cnt):
                    s += [k] * (k - c if c > 0 else -c)
                return int(''.join(map(str, s)))

        return -1  # 无解（本题不会发生，但为了可扩展性保留）

###java

class Solution {
    public int nextBeautifulNumber(int n) {
        // 补一个前导零，方便处理答案十进制串比 n 的十进制串长的情况
        char[] s = ("0" + n).toCharArray();
        int m = s.length;

        final int MX = 10;
        int[] cnt = new int[MX];
        for (int i = 1; i < m; i++) {
            cnt[s[i] - '0']++;
        }

        // 从右往左尝试
        for (int i = m - 1; i >= 0; i--) {
            if (i > 0) {
                cnt[s[i] - '0']--; // 撤销
            }

            // 增大 s[i] 为 j
            for (int j = s[i] - '0' + 1; j < MX; j++) {
                cnt[j]++;

                // 后面 [i+1, m-1] 需要补满 0 < cnt[k] < k 的数字 k，剩余数位可以随便填
                int free = m - 1 - i; // 统计可以随便填的数位个数
                for (int k = 0; k < MX; k++) {
                    int c = cnt[k];
                    if (k < c) { // 不合法
                        free = -1;
                        break;
                    }
                    if (c > 0) {
                        free -= k - c;
                    }
                }
                if (free < 0) { // 不合法，继续枚举
                    cnt[j]--;
                    continue;
                }

                // 对于可以随便填的数位，计算字典序最小的填法
                List<Integer> a = new ArrayList<>();
                for (int k = 1; k < MX; k++) {
                    if (cnt[k] == 0) {
                        a.add(k);
                    }
                }

                List<Integer> missing = zeroOneKnapsack(a, free);
                if (missing == null) { // 无解，继续枚举
                    cnt[j]--;
                    continue;
                }

                for (int v : missing) {
                    cnt[v] = -v; // 用负数表示可以随便填的数
                }

                StringBuilder ans = new StringBuilder("0" + n);
                ans.setCharAt(i, (char) ('0' + j));
                ans.setLength(i + 1);
                for (int k = 1; k < MX; k++) {
                    int c = cnt[k];
                    ans.repeat('0' + k, c > 0 ? k - c : -c);
                }
                return Integer.parseInt(ans.toString());
            }
        }
        return -1; // 无解（本题不会发生，但为了可扩展性保留）
    }

    // 从 a 中选一个字典序最小的、元素和等于 target 的子序列
    // a 已经从小到大排序
    // 无解返回 null
    private List<Integer> zeroOneKnapsack(List<Integer> a, int target) {
        int n = a.size();
        boolean[][] f = new boolean[n + 1][target + 1];
        f[n][0] = true;

        // 倒着 DP，这样后面可以正着（从小到大）选
        for (int i = n - 1; i >= 0; i--) {
            int v = a.get(i);
            for (int j = 0; j <= target; j++) {
                if (j < v) {
                    f[i][j] = f[i + 1][j];
                } else {
                    f[i][j] = f[i + 1][j] || f[i + 1][j - v];
                }
            }
        }

        if (!f[0][target]) {
            return null;
        }

        List<Integer> ans = new ArrayList<>();
        int j = target;
        for (int i = 0; i < n; i++) {
            int v = a.get(i);
            if (j >= v && f[i + 1][j - v]) {
                ans.add(v);
                j -= v;
            }
        }
        return ans;
    }
}

###cpp

class Solution {
    // 从 a 中选一个字典序最小的、元素和等于 target 的子序列
    // a 已经从小到大排序
    // 无解返回 {} 和 false
    pair<vector<int>, bool> zeroOneKnapsack(vector<int>& a, int target) {
        int n = a.size();
        vector f(n + 1, vector<int8_t>(target + 1)); 
        f[n][0] = true;

        // 倒着 DP，这样后面可以正着（从小到大）选
        for (int i = n - 1; i >= 0; i--) {
            int v = a[i];
            for (int j = 0; j <= target; j++) {
                if (j < v) {
                    f[i][j] = f[i + 1][j];
                } else {
                    f[i][j] = f[i + 1][j] || f[i + 1][j - v];
                }
            }
        }

        if (!f[0][target]) {
            return {};
        }

        vector<int> ans;
        int j = target;
        for (int i = 0; i < n; i++) {
            int v = a[i];
            if (j >= v && f[i + 1][j - v]) {
                ans.push_back(v);
                j -= v;
            }
        }
        return {ans, true};
    }

public:
    int nextBeautifulNumber(int n) {
        // 补一个前导零，方便处理答案十进制串比 n 的十进制串长的情况
        string s = "0" + to_string(n);
        int m = s.size();

        constexpr int MX = 10;
        int cnt[MX]{};
        for (int i = 1; i < m; i++) {
            cnt[s[i] - '0']++;
        }

        // 从右往左尝试
        for (int i = m - 1; i >= 0; i--) {
            if (i > 0) {
                cnt[s[i] - '0']--; // 撤销
            }

            // 增大 s[i] 为 j
            for (int j = s[i] - '0' + 1; j < MX; j++) {
                cnt[j]++;

                // 后面 [i+1, m-1] 需要补满 0 < cnt[k] < k 的数字 k，剩余数位可以随便填
                int free = m - 1 - i; // 统计可以随便填的数位个数
                for (int k = 0; k < MX; k++) {
                    int c = cnt[k];
                    if (k < c) { // 不合法
                        free = -1;
                        break;
                    }
                    if (c > 0) {
                        free -= k - c;
                    }
                }
                if (free < 0) { // 不合法，继续枚举
                    cnt[j]--;
                    continue;
                }

                // 对于可以随便填的数位，计算字典序最小的填法
                vector<int> a;
                for (int k = 1; k < MX; k++) {
                    if (cnt[k] == 0) {
                        a.push_back(k);
                    }
                }
                auto [missing, ok] = zeroOneKnapsack(a, free);
                if (!ok) { // 无解，继续枚举
                    cnt[j]--;
                    continue;
                }

                for (int v : missing) {
                    cnt[v] = -v; // 用负数表示可以随便填的数
                }

                s[i] = '0' + j;
                s.resize(i + 1);
                for (int k = 1; k < MX; k++) {
                    int c = cnt[k];
                    c = c > 0 ? k - c : -c;
                    s += string(c, '0' + k);
                }
                return stoi(s);
            }
        }
        return -1; // 无解（本题不会发生，但为了可扩展性保留）
    }
};

###go

// 从 a 中选一个字典序最小的、元素和等于 target 的子序列
// a 已经从小到大排序
// 无解返回 nil
func zeroOneKnapsack(a []int, target int) []int {
n := len(a)
f := make([][]bool, n+1)
for i := range f {
f[i] = make([]bool, target+1)
}
f[n][0] = true

// 倒着 DP，这样后面可以正着（从小到大）选
for i := n - 1; i >= 0; i-- {
v := a[i]
for j := range f[i] {
if j < v {
f[i][j] = f[i+1][j]
} else {
f[i][j] = f[i+1][j] || f[i+1][j-v]
}
}
}

if !f[0][target] {
return nil
}

ans := []int{}
j := target
for i, v := range a {
if j >= v && f[i+1][j-v] {
ans = append(ans, v)
j -= v
}
}
return ans
}

func nextBeautifulNumber(n int) int {
// 补一个前导零，方便处理答案十进制比 n 的十进制长的情况
s := "0" + strconv.Itoa(n)
m := len(s)

const mx = 10
cnt := make([]int, mx)
for i := 1; i < m; i++ {
cnt[s[i]-'0']++
}

// 从右往左尝试
for i := m - 1; i >= 0; i-- {
if i > 0 {
cnt[s[i]-'0']-- // 撤销
}

// 增大 s[i] 为 j
for j := s[i] - '0' + 1; j < mx; j++ {
cnt[j]++

// 后面 [i+1, m-1] 需要补满 0 < cnt[k] < k 的数字 k，剩余数位可以随便填
free := m - 1 - i // 统计可以随便填的数位个数
for k, c := range cnt {
if k < c { // 不合法
free = -1
break
}
if c > 0 {
free -= k - c
}
}
if free < 0 { // 不合法，继续枚举
cnt[j]--
continue
}

// 对于可以随便填的数位，计算字典序最小的填法
a := []int{}
for k := 1; k < mx; k++ {
if cnt[k] == 0 {
a = append(a, k)
}
}
missing := zeroOneKnapsack(a, free)
if missing == nil { // 无解，继续枚举
cnt[j]--
continue
}

for _, v := range missing {
cnt[v] = -v // 用负数表示可以随便填的数
}

t := []byte(s[:i+1])
t[i] = '0' + byte(j)
for k, c := range cnt {
if c > 0 {
c = k - c
} else {
c = -c
}
d := []byte{'0' + byte(k)}
t = append(t, bytes.Repeat(d, c)...)
}
ans, _ := strconv.Atoi(string(t))
return ans
}
}
return -1 // 无解（本题不会发生，但为了可扩展性保留）
}

复杂度分析

时间复杂度：$\mathcal{O}(D^2\log^2 n)$，其中 $D=10$，$\log n$ 是 $n$ 的十进制长度。枚举 $\mathcal{O}(D\log n)$ 种把 $s[i]$ 增大的情况，每次需要 $\mathcal{O}(D\log n)$ 的时间计算 0-1 背包。
空间复杂度：$\mathcal{O}(D\log n)$。

相似题目

专题训练

见下面贪心题单的「§3.1 字典序最小/最大」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

非暴力做法（Python/Java/C++/Go）

LeetCode 每日一题题解

作者 endlesscheng

2025年2月23日 12:21

本题和周赛第三题是一样的，请看我的题解。

LeetCode 每日一题题解
两种方法：差分/三指针+双指针（Python/Java/C++/Go）endlesscheng
2024年11月10日 09:14

两种方法：差分/三指针+双指针（Python/Java/C++/Go）

LeetCode 每日一题题解

作者 endlesscheng

2024年11月10日 09:14

方法一：差分

前置知识：差分原理讲解。

设 $x = \textit{nums}[i]$。根据题意，$x$ 可以变成 $[x-k,x+k]$ 中的整数。

题目让我们计算操作后，最多有多少个数相同。

例如 $\textit{nums}=[2,4]$，$k=1$，$\textit{numOperations}=2$。$2$ 可以变成 $[1,3]$ 中的整数，$4$ 可以变成 $[3,5]$ 中的整数。$2$ 和 $4$ 都可以变成 $3$，所以答案是 $2$。

一般地，$x$ 可以变成 $[x-k,x+k]$ 中的整数，我们可以把 $[x-k,x+k]$ 中的每个整数的出现次数都加一，然后统计出现次数的最大值。这可以用差分实现。

计算差分的前缀和。设有 $\textit{sumD}$ 个数可以变成 $y$。

如果 $y$ 不在 $\textit{nums}$ 中，那么 $y$ 的最大出现次数不能超过 $\textit{numOperations}$，与 $\textit{sumD}$ 取最小值，得 $\min(\textit{sumD}, \textit{numOperations})$。

如果 $y$ 在 $\textit{nums}$ 中，且出现了 $\textit{cnt}$ 次，那么有 $\textit{sumD}-\textit{cnt}$ 个其他元素（不等于 $y$ 的数）可以变成 $y$，但这不能超过 $\textit{numOperations}$，所以有

$$
\min(\textit{sumD}-\textit{cnt}, \textit{numOperations})
$$

个其他元素可以变成 $y$，再加上 $y$ 自身出现的次数 $\textit{cnt}$，得到 $y$ 的最大频率为

$$
\textit{cnt} + \min(\textit{sumD}-\textit{cnt}, \textit{numOperations}) = \min(\textit{sumD}, \textit{cnt}+\textit{numOperations})
$$

注意上式兼容 $y$ 不在 $\textit{nums}$ 中的情况，此时 $\textit{cnt}=0$。

本题视频讲解，欢迎点赞关注~

答疑

问：为什么代码只考虑在 $\textit{diff}$ 和 $\textit{nums}$ 中的数字？

答：比如 $x$ 在 $\textit{diff}$ 中，$x+1,x+2,\ldots$ 不在 $\textit{diff}$ 中，那么 $x+1,x+2,\ldots$ 的 $\textit{sumD}$ 和 $\textit{x}$ 的是一样的，无需重复计算。此外，要想算出比 $\min(\textit{sumD}, \textit{cnt}+\textit{numOperations})$ 更大的数，要么 $\textit{sumD}$ 变大，要么 $\textit{cnt}$ 变大。「变大」时的 $x$ 必然在 $\textit{diff}$ 或 $\textit{nums}$ 中。

###py

class Solution:
    def maxFrequency(self, nums: List[int], k: int, numOperations: int) -> int:
        cnt = defaultdict(int)
        diff = defaultdict(int)
        for x in nums:
            cnt[x] += 1
            diff[x]  # 把 x 插入 diff，以保证下面能遍历到 x
            diff[x - k] += 1  # 把 [x-k,x+k] 中的每个整数的出现次数都加一
            diff[x + k + 1] -= 1

        ans = sum_d = 0
        for x, d in sorted(diff.items()):
            sum_d += d
            ans = max(ans, min(sum_d, cnt[x] + numOperations))
        return ans

###java

class Solution {
    int maxFrequency(int[] nums, int k, int numOperations) {
        Map<Integer, Integer> cnt = new HashMap<>();
        Map<Integer, Integer> diff = new TreeMap<>();
        for (int x : nums) {
            cnt.merge(x, 1, Integer::sum); // cnt[x]++
            diff.putIfAbsent(x, 0); // 把 x 插入 diff，以保证下面能遍历到 x
            // 把 [x-k, x+k] 中的每个整数的出现次数都加一
            diff.merge(x - k, 1, Integer::sum); // diff[x-k]++
            diff.merge(x + k + 1, -1, Integer::sum); // diff[x+k+1]--
        }

        int ans = 0;
        int sumD = 0;
        for (Map.Entry<Integer, Integer> e : diff.entrySet()) {
            sumD += e.getValue();
            ans = Math.max(ans, Math.min(sumD, cnt.getOrDefault(e.getKey(), 0) + numOperations));
        }
        return ans;
    }
}

###cpp

class Solution {
public:
    int maxFrequency(vector<int>& nums, int k, int numOperations) {
        unordered_map<int, int> cnt;
        map<int, int> diff;
        for (int x : nums) {
            cnt[x]++;
            diff[x]; // 把 x 插入 diff，以保证下面能遍历到 x
            diff[x - k]++; // 把 [x-k, x+k] 中的每个整数的出现次数都加一
            diff[x + k + 1]--;
        }

        int ans = 0, sum_d = 0;
        for (auto& [x, d] : diff) {
            sum_d += d;
            ans = max(ans, min(sum_d, cnt[x] + numOperations));
        }
        return ans;
    }
};

###go

func maxFrequency(nums []int, k, numOperations int) (ans int) {
cnt := map[int]int{}
diff := map[int]int{}
for _, x := range nums {
cnt[x]++
diff[x] += 0 // 把 x 插入 diff，以保证下面能遍历到 x
diff[x-k]++  // 把 [x-k, x+k] 中的每个整数的出现次数都加一
diff[x+k+1]--
}

sumD := 0
for _, x := range slices.Sorted(maps.Keys(diff)) {
sumD += diff[x]
ans = max(ans, min(sumD, cnt[x]+numOperations))
}
return
}

复杂度分析

时间复杂度：$\mathcal{O}(n\log n)$，其中 $n$ 是 $\textit{nums}$ 的长度。
空间复杂度：$\mathcal{O}(n)$。

方法二：同向三指针 + 同向双指针

核心思路

计算有多少个数能变成 $x$，其中 $x = \textit{nums}[i]$。用同向三指针实现。
计算有多少个数能变成 $x$，其中 $x$ 不一定在 $\textit{nums}$ 中。用同向双指针实现。

同向三指针

把 $\textit{nums}$ 从小到大排序。

遍历 $\textit{nums}$。设 $x=\textit{nums}[i]$，计算元素值在 $[x-k,x+k]$ 中的元素个数，这些元素都可以变成 $x$。

遍历的同时，维护左指针 $\textit{left}$，它是最小的满足

$$
\textit{nums}[\textit{left}] \ge x - k
$$

的下标。

遍历的同时，维护右指针 $\textit{right}$，它是最小的满足

$$
\textit{nums}[\textit{right}] > x + k
$$

的下标。如果不存在，则 $\textit{right}=n$。

下标在左闭右开区间 $[\textit{left},\textit{right})$ 中的元素，都可以变成 $x$。这有

$$
\textit{sumD} = \textit{right} - \textit{left}
$$

个。

遍历的同时，求出 $x$ 有 $\textit{cnt}$ 个。然后用方法一的公式，更新答案的最大值。

同向双指针

同向三指针没有考虑「变成不在 $\textit{nums}$ 中的数」这种情况。

然而，不在 $\textit{nums}$ 中的数太多了！怎么减少计算量？

假设都变成 $y$，那么只有 $[y-k,y+k]$ 中的数能变成 $y$。

把 $\textit{nums}[i]$ 视作一维坐标轴上的点，想象一个窗口 $[y-k,y+k]$ 在不断地向右滑动，什么时候窗口内的元素个数才会变大？

如果我们从 $[y-k-1,y+k-1]$ 移动到 $[y-k,y+k]$，且 $y+k$ 不在 $\textit{nums}$ 中，此时窗口内的元素个数不会变大，甚至因为左端点收缩了，元素个数可能会变小。
所以，只有当 $y+k$ 恰好在 $\textit{nums}$ 中时，窗口内的元素个数才可能会变大。
结论：我们只需考虑 $y+k$ 在 $\textit{nums}$ 中时的 $y$！

于是，枚举 $x=\textit{nums}[\textit{right}]$，计算元素值在 $[x-2k,x]$ 中的元素个数，这些元素都可以变成同一个数 $y=x-k$。

左指针 $\textit{left}$ 是最小的满足

$$
\textit{nums}[\textit{left}] \ge x-2k
$$

的下标。

计算好 $\textit{left}$ 后，下标在 $[\textit{left}, \textit{right}]$ 中的数可以变成一样的，这有

$$
\textit{right} - \textit{left} + 1
$$

个。注意上式不能超过 $\textit{numOperations}$。

小优化

由于同向双指针算出的结果不超过 $\textit{numOperations}$，所以当同向三指针计算完毕后，如果发现答案已经 $\ge \textit{numOperations}$，那么无需计算同向双指针。

###py

class Solution:
    def maxFrequency(self, nums: List[int], k: int, numOperations: int) -> int:
        nums.sort()

        n = len(nums)
        ans = cnt = left = right = 0
        for i, x in enumerate(nums):
            cnt += 1
            # 循环直到连续相同段的末尾，这样可以统计出 x 的出现次数
            if i < n - 1 and x == nums[i + 1]:
                continue
            while nums[left] < x - k:
                left += 1
            while right < n and nums[right] <= x + k:
                right += 1
            ans = max(ans, min(right - left, cnt + numOperations))
            cnt = 0

        if ans >= numOperations:
            return ans

        left = 0
        for right, x in enumerate(nums):
            while nums[left] < x - k * 2:
                left += 1
            ans = max(ans, right - left + 1)
        return min(ans, numOperations)  # 最后和 numOperations 取最小值

###java

class Solution {
    public int maxFrequency(int[] nums, int k, int numOperations) {
        Arrays.sort(nums);

        int n = nums.length;
        int ans = 0;
        int cnt = 0;
        int left = 0;
        int right = 0;
        for (int i = 0; i < n; i++) {
            int x = nums[i];
            cnt++;
            // 循环直到连续相同段的末尾，这样可以统计出 x 的出现次数
            if (i < n - 1 && x == nums[i + 1]) {
                continue;
            }
            while (nums[left] < x - k) {
                left++;
            }
            while (right < n && nums[right] <= x + k) {
                right++;
            }
            ans = Math.max(ans, Math.min(right - left, cnt + numOperations));
            cnt = 0;
        }

        if (ans >= numOperations) {
            return ans;
        }

        left = 0;
        for (right = 0; right < n; right++) {
            int x = nums[right];
            while (nums[left] < x - k * 2) {
                left++;
            }
            ans = Math.max(ans, right - left + 1);
        }
        return Math.min(ans, numOperations); // 最后和 numOperations 取最小值
    }
}

###cpp

class Solution {
public:
    int maxFrequency(vector<int>& nums, int k, int numOperations) {
        ranges::sort(nums);

        int n = nums.size();
        int ans = 0, cnt = 0, left = 0, right = 0;
        for (int i = 0; i < n; i++) {
            int x = nums[i];
            cnt++;
            // 循环直到连续相同段的末尾，这样可以统计出 x 的出现次数
            if (i < n - 1 && x == nums[i + 1]) {
                continue;
            }
            while (nums[left] < x - k) {
                left++;
            }
            while (right < n && nums[right] <= x + k) {
                right++;
            }
            ans = max(ans, min(right - left, cnt + numOperations));
            cnt = 0;
        }

        if (ans >= numOperations) {
            return ans;
        }

        left = 0;
        for (int right = 0; right < n; right++) {
            int x = nums[right];
            while (nums[left] < x - k * 2) {
                left++;
            }
            ans = max(ans, right - left + 1);
        }
        return min(ans, numOperations); // 最后和 numOperations 取最小值
    }
};

###go

func maxFrequency(nums []int, k, numOperations int) (ans int) {
slices.Sort(nums)

n := len(nums)
var cnt, left, right int
for i, x := range nums {
cnt++
// 循环直到连续相同段的末尾，这样可以统计出 x 的出现次数
if i < n-1 && x == nums[i+1] {
continue
}
for nums[left] < x-k {
left++
}
for right < n && nums[right] <= x+k {
right++
}
ans = max(ans, min(right-left, cnt+numOperations))
cnt = 0
}

if ans >= numOperations {
return ans
}

left = 0
for right, x := range nums {
for nums[left] < x-k*2 {
left++
}
ans = max(ans, right-left+1)
}
return min(ans, numOperations) // 最后和 numOperations 取最小值
}

也可以把两个 for 循环合起来。

###py

class Solution:
    def maxFrequency(self, nums: List[int], k: int, numOperations: int) -> int:
        nums.sort()

        n = len(nums)
        ans = cnt = left = right = left2 = 0
        for i, x in enumerate(nums):
            while nums[left2] < x - k * 2:
                left2 += 1
            ans = max(ans, min(i - left2 + 1, numOperations))

            cnt += 1
            # 循环直到连续相同段的末尾，这样可以统计出 x 的出现次数
            if i < n - 1 and x == nums[i + 1]:
                continue
            while nums[left] < x - k:
                left += 1
            while right < n and nums[right] <= x + k:
                right += 1
            ans = max(ans, min(right - left, cnt + numOperations))
            cnt = 0

        return ans

###java

class Solution {
    public int maxFrequency(int[] nums, int k, int numOperations) {
        Arrays.sort(nums);

        int n = nums.length;
        int ans = 0;
        int cnt = 0;
        int left = 0;
        int right = 0;
        int left2 = 0;
        for (int i = 0; i < n; i++) {
            int x = nums[i];
            while (nums[left2] < x - k * 2) {
                left2++;
            }
            ans = Math.max(ans, Math.min(i - left2 + 1, numOperations));

            cnt++;
            // 循环直到连续相同段的末尾，这样可以统计出 x 的出现次数
            if (i < n - 1 && x == nums[i + 1]) {
                continue;
            }
            while (nums[left] < x - k) {
                left++;
            }
            while (right < n && nums[right] <= x + k) {
                right++;
            }
            ans = Math.max(ans, Math.min(right - left, cnt + numOperations));
            cnt = 0;
        }

        return ans;
    }
}

###cpp

class Solution {
public:
    int maxFrequency(vector<int>& nums, int k, int numOperations) {
        ranges::sort(nums);

        int n = nums.size();
        int ans = 0, cnt = 0, left = 0, right = 0, left2 = 0;
        for (int i = 0; i < n; i++) {
            int x = nums[i];
            while (nums[left2] < x - k * 2) {
                left2++;
            }
            ans = max(ans, min(i - left2 + 1, numOperations));

            cnt++;
            // 循环直到连续相同段的末尾，这样可以统计出 x 的出现次数
            if (i < n - 1 && x == nums[i + 1]) {
                continue;
            }
            while (nums[left] < x - k) {
                left++;
            }
            while (right < n && nums[right] <= x + k) {
                right++;
            }
            ans = max(ans, min(right - left, cnt + numOperations));
            cnt = 0;
        }

        return ans;
    }
};

###go

func maxFrequency(nums []int, k, numOperations int) (ans int) {
slices.Sort(nums)

n := len(nums)
var cnt, left, right, left2 int
for i, x := range nums {
for nums[left2] < x-k*2 {
left2++
}
ans = max(ans, min(i-left2+1, numOperations))

cnt++
// 循环直到连续相同段的末尾，这样可以统计出 x 的出现次数
if i < n-1 && x == nums[i+1] {
continue
}
for nums[left] < x-k {
left++
}
for right < n && nums[right] <= x+k {
right++
}
ans = max(ans, min(right-left, cnt+numOperations))
cnt = 0
}

return
}

复杂度分析

时间复杂度：$\mathcal{O}(n\log n)$，其中 $n$ 是 $\textit{nums}$ 的长度。瓶颈在排序上。
空间复杂度：$\mathcal{O}(1)$。忽略排序的栈开销。

专题训练

下面数据结构题单的「§2.1 一维差分」。
下面双指针题单的「§3.2 同向双指针」和「五、三指针」。

分类题单

如何科学刷题？

我的题解精选（已分类）

LeetCode 每日一题题解
枚举轮转到最左边的下标 + 裴蜀定理（Python/Java/C++/C/Go/JS/Rust）endlesscheng
2025年10月3日 10:35

枚举轮转到最左边的下标 + 裴蜀定理（Python/Java/C++/C/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2025年10月3日 10:35

注：题干中的「数字一旦超过 $9$ 就会变成 $0$」的意思是，数字 $x$ 加上 $a$ 后，会变成 $(x+a)\bmod 10$。

不轮转

从特殊到一般，先考虑只累加、不轮转的情况。

为了让字典序尽量小，第一个奇数下标 $1$ 上的数字 $s_1$，越小越好。一旦我们确定了 $s_1$ 的最终值，就确定了一共累加的值。由于所有奇数下标都要累加同一个数，所以也就确定了其余奇数下标的值。

那么，$s_1$ 最小可以是多少？可以是 $0$ 吗？

比如 $s_1 = 5$，$a=2$。不断累加 $a$，$s_1$ 的变化情况如下：

$$
5\to 7\to 9\to 1\to 3\to 5\to \cdots
$$

这种情况 $s_1$ 只能是奇数，最小是 $1$。

而如果 $s_1 = 5$，$a=3$，$s_1$ 的变化情况如下：

$$
5\to 8\to 1\to 4\to 7\to 0\to 3\to 6\to 9\to 2\to 5 \cdots
$$

这种情况 $s_1$ 可以变成 $[0,9]$ 中的任意整数，最小是 $0$。

一般地，设累加操作执行了 $k\ (k\ge 0)$ 次，那么 $s_1$ 变成

$$
r = (s_1 + ak)\bmod 10
$$

也就是 $s_1 + ak$ 减去若干个 $10$ 等于 $r$，即

$$
s_1 + ak - 10q = r
$$

变形得

$$
ak - 10q = r-s_1
$$

裴蜀定理指出，该方程有解，当且仅当 $r-s_1$ 是 $g = \gcd(a,10)$ 的倍数，即

$$
r \equiv s_1 \pmod g
$$

其中 $\equiv$ 是同余符号，详细解释见模运算的世界：当加减乘除遇上取模。

上式表明，$s_1$ 通过累加操作变成的数，必须与 $s_1$ 关于模 $g$ 同余，所以 $s_1$ 可以变成的最小值为

$$
s_1\bmod g
$$

从 $s_1$ 到 $s_1\bmod g$，一共要累加的值为

$$
s_1\bmod g - s_1 + 10
$$

其中 $+10$ 保证减法结果非负。

枚举轮转到最左边的下标

例如 $s=\texttt{012345}$，$b=4$，执行轮转操作，得到

$$
\texttt{012345}\to\texttt{234501}\to\texttt{450123}\to\texttt{012345}\to\cdots
$$

只有 $s_0,s_2,s_4$ 可以轮转到最左边。

类似上文的思路，根据裴蜀定理，可以轮转到最左边的下标，必须是 $\textit{step} = \gcd(b,n)$ 的倍数，其中 $n$ 是 $s$ 的长度。

枚举 $i = 0,\textit{step},2\cdot \textit{step}, 3\cdot \textit{step},\dots$ 作为轮转到最左边的下标。

分类讨论：

如果 $\gcd(b,n)$ 是偶数，无论如何轮转，我们只能对奇数下标执行累加操作。
如果 $\gcd(b,n)$ 是奇数，轮转一次后，原来的偶数下标变成奇数下标。可以先轮转一次，执行累加，再轮转到我们想要的位置。可以视作我们拥有了「对偶数下标执行累加操作」的能力。

###py

class Solution:
    def findLexSmallestString(self, s: str, a: int, b: int) -> str:
        s = list(map(int, s))
        n = len(s)
        step = gcd(b, n)
        g = gcd(a, 10)
        ans = [inf]

        def modify(start: int) -> None:
            ch = t[start]  # 最靠前的数字，越小越好
            # ch 可以变成的最小值为 ch%g
            # 例如 ch=5，g=2，那么 ch+2+2+2（模 10）后变成 1，不可能变得更小
            # 从 ch 到 ch%g，需要增加 inc（循环中会 %10 保证结果在 [0,9] 中）
            inc = ch % g - ch
            if inc:  # 优化：inc 为 0 时，t[j] 不变，无需执行 for 循环
                for j in range(start, n, 2):
                    t[j] = (t[j] + inc) % 10

        for i in range(0, n, step):      
            t = s[i:] + s[:i]  # 轮转
            modify(1)  # 累加操作（所有奇数下标）
            if step % 2:  # 能对偶数下标执行累加操作
                modify(0)  # 累加操作（所有偶数下标）
            ans = min(ans, t)

        return ''.join(map(str, ans))

###java

class Solution {
    public String findLexSmallestString(String S, int a, int b) {
        char[] s = S.toCharArray();
        int n = s.length;
        char[] t = new char[n];
        int step = gcd(b, n);
        int g = gcd(a, 10);
        String ans = null;

        for (int i = 0; i < n; i += step) {
            // t = s[i,n) + s[0,i)
            System.arraycopy(s, i, t, 0, n - i);
            System.arraycopy(s, 0, t, n - i, i);

            modify(t, 1, g); // 累加操作（所有奇数下标）
            if (step % 2 > 0) { // 能对偶数下标执行累加操作
                modify(t, 0, g); // 累加操作（所有偶数下标）
            }

            String str = new String(t);
            if (ans == null || str.compareTo(ans) < 0) {
                ans = str;
            }
        }

        return ans;
    }

    private void modify(char[] t, int start, int g) {
        int ch = t[start] - '0'; // 最靠前的数字，越小越好
        // ch 可以变成的最小值为 ch%g
        // 例如 ch=5，g=2，那么 ch+2+2+2（模 10）后变成 1，不可能变得更小
        // 从 ch 到 ch%g，需要增加 inc，其中 +10 保证 inc 非负（循环中会 %10 保证结果在 [0,9] 中）
        int inc = ch % g - ch + 10;
        for (int j = start; j < t.length; j += 2) {
            t[j] = (char) ('0' + (t[j] - '0' + inc) % 10);
        }
    }

    private int gcd(int a, int b) {
        while (a != 0) {
            int tmp = a;
            a = b % a;
            b = tmp;
        }
        return b;
    }
}

###cpp

class Solution {
public:
    string findLexSmallestString(string s, int a, int b) {
        int n = s.size();
        int step = gcd(b, n);
        int g = gcd(a, 10);
        string ans;

        for (int i = 0; i < n; i += step) {
            string t = s.substr(i) + s.substr(0, i); // 轮转

            auto modify = [&](int start) -> void {
                int ch = t[start] - '0'; // 最靠前的数字，越小越好
                // ch 可以变成的最小值为 ch%g
                // 例如 ch=5，g=2，那么 ch+2+2+2（模 10）后变成 1，不可能变得更小
                // 从 ch 到 ch%g，需要增加 inc，其中 +10 保证 inc 非负（循环中会 %10 保证结果在 [0,9] 中）
                int inc = ch % g - ch + 10;
                for (int j = start; j < n; j += 2) {
                    t[j] = '0' + (t[j] - '0' + inc) % 10;
                }
            };

            modify(1); // 累加操作（所有奇数下标）
            if (step % 2) { // 能对偶数下标执行累加操作
                modify(0); // 累加操作（所有偶数下标）
            }

            if (ans.empty() || t < ans) {
                ans = move(t);
            }
        }

        return ans;
    }
};

###c

int gcd(int a, int b) {
    while (a) {
        int tmp = a;
        a = b % a;
        b = tmp;
    }
    return b;
}

void modify(char* t, int n, int start, int g) {
    int ch = t[start] - '0'; // 最靠前的数字，越小越好
    // ch 可以变成的最小值为 ch%g
    // 例如 ch=5，g=2，那么 ch+2+2+2（模 10）后变成 1，不可能变得更小
    // 从 ch 到 ch%g，需要增加 inc，其中 +10 保证 inc 非负（循环中会 %10 保证结果在 [0,9] 中）
    int inc = ch % g - ch + 10;
    for (int j = start; j < n; j += 2) {
        t[j] = '0' + (t[j] - '0' + inc) % 10;
    }
}

char* findLexSmallestString(char* s, int a, int b) {
    int n = strlen(s);
    int step = gcd(b, n);
    int g = gcd(a, 10);

    char* ans = malloc((n + 1) * sizeof(char));
    ans[0] = CHAR_MAX;
    ans[1] = '\0';

    char* t = malloc((n + 1) * sizeof(char));
    t[n] = '\0';

    for (int i = 0; i < n; i += step) {
        // t = s[i,n) + s[0,i)
        strncpy(t, s + i, n - i);
        strncpy(t + n - i, s, i);

        modify(t, n, 1, g); // 累加操作（所有奇数下标）
        if (step % 2) { // 能对偶数下标执行累加操作
            modify(t, n, 0, g); // 累加操作（所有偶数下标）
        }

        if (strcmp(t, ans) < 0) {
            strcpy(ans, t);
        }
    }

    free(t);
    return ans;
}

###go

func findLexSmallestString(s string, a int, b int) string {
n := len(s)
step := gcd(b, n)
g := gcd(a, 10)
var ans []byte

for i := 0; i < n; i += step {
t := []byte(s[i:] + s[:i]) // 轮转
modify := func(start int) {
ch := t[start] - '0' // 最靠前的数字，越小越好
// ch 可以变成的最小值为 ch%g
// 例如 ch=5，g=2，那么 ch+2+2+2（模 10）后变成 1，不可能变得更小
// 从 ch 到 ch%g，需要增加 inc，其中 +10 保证 inc 非负（循环中会 %10 保证结果在 [0,9] 中）
inc := ch%byte(g) + 10 - ch
for j := start; j < n; j += 2 {
t[j] = '0' + (t[j]-'0'+inc)%10
}
}
modify(1) // 累加操作（所有奇数下标）
if step%2 > 0 { // 能对偶数下标执行累加操作
modify(0) // 累加操作（所有偶数下标）
}
if ans == nil || bytes.Compare(t, ans) < 0 {
ans = t
}
}

return string(ans)
}

func gcd(a, b int) int {
for a != 0 {
a, b = b%a, a
}
return b
}

###js

var findLexSmallestString = function(s, a, b) {
    const arr = s.split('').map(ch => parseInt(ch));
    const n = arr.length;
    const step = gcd(b, n);
    const g = gcd(a, 10);
    let ans = null;

    function modify(t, start) {
        const ch = t[start]; // 最靠前的数字，越小越好
        // ch 可以变成的最小值为 ch%g
        // 例如 ch=5，g=2，那么 ch+2+2+2（模 10）后变成 1，不可能变得更小
        // 从 ch 到 ch%g，需要增加 inc（循环中会 %10 保证结果在 [0,9] 中）
        let inc = ch % g - ch + 10;
        if (inc === 0) { // 优化：inc 为 0 时，t[j] 不变，无需执行 for 循环
            return;
        }
        for (let j = start; j < n; j += 2) {
            t[j] = (t[j] + inc) % 10;
        }
    }

    for (let i = 0; i < n; i += step) {
        const t = arr.slice(i).concat(arr.slice(0, i)); // 轮转
        modify(t, 1); // 累加操作（所有奇数下标）
        if (step % 2) { // 能对偶数下标执行累加操作
            modify(t, 0); // 累加操作（所有偶数下标）
        }
        if (ans === null || compareArray(t, ans) < 0) {
            ans = t;
        }
    }

    return ans.join('');
};

function gcd(a, b) {
    while (a) {
        [a, b] = [b % a, a];
}
return b;
}

function compareArray(a, b) {
    const n = a.length;
    for (let i = 0; i < n; i++) {
        if (a[i] !== b[i]) {
            return a[i] - b[i];
        }
    }
    return 0;
}

###rust

impl Solution {
    pub fn find_lex_smallest_string(s: String, a: i32, b: i32) -> String {
        let n = s.len();
        let step = gcd(b, n as i32) as usize;
        let g = gcd(a, 10) as u8;
        let mut ans = vec![u8::MAX];

        let modify = |t: &mut [u8], start: usize| {
            let ch = t[start] - b'0'; // 最靠前的数字，越小越好
            // ch 可以变成的最小值为 ch%g
            // 例如 ch=5，g=2，那么 ch+2+2+2（模 10）后变成 1，不可能变得更小
            // 从 ch 到 ch%g，需要增加 inc，其中 +10 保证 inc 非负（循环中会 %10 保证结果在 [0,9] 中）
            let inc = ch % g + 10 - ch;
            for j in (start..n).step_by(2) {
                t[j] = b'0' + (t[j] - b'0' + inc) % 10;
            }
        };

        for i in (0..n).step_by(step) {
            let mut t = format!("{}{}", &s[i..], &s[..i]).into_bytes(); // 轮转
            modify(&mut t, 1); // 累加操作（所有奇数下标）
            if step % 2 != 0 { // 能对偶数下标执行累加操作
                modify(&mut t, 0); // 累加操作（所有偶数下标）
            }
            ans = ans.min(t);
        }

        unsafe { String::from_utf8_unchecked(ans) }
    }
}

fn gcd(mut a: i32, mut b: i32) -> i32 {
    while a != 0 {
        (a, b) = (b % a, a);
    }
    b
}

复杂度分析

时间复杂度：$\mathcal{O}\left(\dfrac{n^2}{\gcd(b,n)}\right)$，其中 $n$ 是 $s$ 的长度。
空间复杂度：$\mathcal{O}(n)$。

注：还可以枚举奇数下标累加值，枚举偶数下标累加值，然后用类似最小表示法的思想，计算在固定累加值的情况下，轮转后的最小字典序。时间复杂度 $\mathcal{O}(D^2n)$，$D=10$。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

从小到大贪心（Python/Java/C++/Go）

LeetCode 每日一题题解

作者 endlesscheng

2024年12月22日 12:05

一个现实中的例子：

军训的某一天，同学们在操场上。现在教官吹响了口哨，同学们集合，排成一排。对于最靠左的同学 $A$，他需要尽量往左移动，给其他同学腾出位置。$A$ 旁边的同学，可以紧挨着 $A$。依此类推。

把 $\textit{nums}$ 视作 $n$ 个同学在一维数轴中的位置，从最左边的同学（$\textit{nums}$ 的最小值）开始思考。

设最左边的同学的位置为 $a$，他尽量往左移，位置变成 $a-k$。

$\textit{nums}$ 的次小值 $b$ 呢？这位同学也尽量往左移：

比如 $a=4,b=6,k=3$，那么 $a$ 变成 $a-k=1$，$b$ 变成 $b-k=3$。
比如 $a=4,b=4,k=3$，那么 $a$ 变成 $a'=a-k=1$，$b$ 变成 $b-k=1$ 就和 $a'$ 一样了，可以稍微大一点（紧挨着 $a'$），把 $b$ 变成 $a'+1=2$。

一般地，$b$ 变成

$$
\max(b-k,a'+1)
$$

但这不能超过 $b+k$，所以最终要变成

$$
\min(\max(b-k,a'+1),b+k)
$$

相当于让 $a'+1$ 落在 $[b-k,b+k]$ 中，若超出范围则修正。

第三小的数也同理，通过前一个数可以算出当前元素能变成多少。

最后答案为 $\textit{nums}$ 中的不同元素个数。我们可以在修改的同时统计，如果发现当前元素修改后的值，比上一个元素修改后的值大，那么答案加一。

为方便计算，把 $\textit{nums}$ 从小到大排序。排序后，从左到右遍历数组，就相当于从最左边的人开始计算了。

本题视频讲解，欢迎点赞关注~

###py

class Solution:
    def maxDistinctElements(self, nums: List[int], k: int) -> int:
        nums.sort()
        ans = 0
        pre = -inf  # 记录每个人左边的人的位置
        for x in nums:
            x = min(max(x - k, pre + 1), x + k)
            if x > pre:
                ans += 1
                pre = x
        return ans

###java

class Solution {
    public int maxDistinctElements(int[] nums, int k) {
        Arrays.sort(nums);
        int ans = 0;
        int pre = Integer.MIN_VALUE; // 记录每个人左边的人的位置
        for (int x : nums) {
            x = Math.min(Math.max(x - k, pre + 1), x + k);
            if (x > pre) {
                ans++;
                pre = x;
            }
        }
        return ans;
    }
}

###cpp

class Solution {
public:
    int maxDistinctElements(vector<int>& nums, int k) {
        ranges::sort(nums);
        int ans = 0;
        int pre = INT_MIN; // 记录每个人左边的人的位置
        for (int x : nums) {
            x = clamp(pre + 1, x - k, x + k); // min(max(x - k, pre + 1), x + k)
            if (x > pre) {
                ans++;
                pre = x;
            }
        }
        return ans;
    }
};

###go

func maxDistinctElements(nums []int, k int) (ans int) {
slices.Sort(nums)
pre := math.MinInt // 记录每个人左边的人的位置
for _, x := range nums {
x = min(max(x-k, pre+1), x+k)
if x > pre {
ans++
pre = x
}
}
return
}

优化

什么情况下，可以直接返回 $n$？

先考虑 $\textit{nums}$ 所有元素都相同的情况（同学们都挤在一起）。我们可以把元素 $x$ 变成 $[x-k,x+k]$ 中的整数，这一共有 $2k+1$ 个。如果 $2k+1 \ge n$，就可以让所有元素互不相同。

如果 $\textit{nums}$ 有不同元素，当 $2k+1 \ge n$ 时，更加可以让所有元素互不相同。

所以只要 $2k+1 \ge n$，就可以直接返回 $n$。

###py

class Solution:
    def maxDistinctElements(self, nums: List[int], k: int) -> int:
        if k * 2 + 1 >= len(nums):
            return len(nums)

        nums.sort()
        ans = 0
        pre = -inf  # 记录每个人左边的人的位置
        for x in nums:
            x = min(max(x - k, pre + 1), x + k)
            if x > pre:
                ans += 1
                pre = x
        return ans

###java

class Solution {
    public int maxDistinctElements(int[] nums, int k) {
        int n = nums.length;
        if (k * 2 + 1 >= n) {
            return n;
        }

        Arrays.sort(nums);
        int ans = 0;
        int pre = Integer.MIN_VALUE; // 记录每个人左边的人的位置
        for (int x : nums) {
            x = Math.min(Math.max(x - k, pre + 1), x + k);
            if (x > pre) {
                ans++;
                pre = x;
            }
        }
        return ans;
    }
}

###cpp

class Solution {
public:
    int maxDistinctElements(vector<int>& nums, int k) {
        int n = nums.size();
        if (k * 2 + 1 >= n) {
            return n;
        }

        ranges::sort(nums);
        int ans = 0;
        int pre = INT_MIN; // 记录每个人左边的人的位置
        for (int x : nums) {
            x = clamp(pre + 1, x - k, x + k); // min(max(x - k, pre + 1), x + k)
            if (x > pre) {
                ans++;
                pre = x;
            }
        }
        return ans;
    }
};

###go

func maxDistinctElements(nums []int, k int) (ans int) {
n := len(nums)
if k*2+1 >= n {
return n
}

slices.Sort(nums)
pre := math.MinInt // 记录每个人左边的人的位置
for _, x := range nums {
x = min(max(x-k, pre+1), x+k)
if x > pre {
ans++
pre = x
}
}
return
}

复杂度分析

时间复杂度：$\mathcal{O}(n\log n)$，其中 $n$ 是 $\textit{nums}$ 的长度。瓶颈在排序上。
空间复杂度：$\mathcal{O}(1)$。忽略排序的栈开销。

专题训练

见下面贪心题单的「§1.1 从最小/最大开始贪心」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
同余分组（Python/Java/C++/C/Go/JS/Rust）endlesscheng
2023年3月19日 12:06

同余分组（Python/Java/C++/C/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2023年3月19日 12:06

下文记 $m=\textit{value}$。

由于同一个数可以加减任意倍的 $m$，我们可以先把每个 $\textit{nums}[i]$ 变成与 $\textit{nums}[i]$ 关于模 $m$ 同余的最小非负整数，以备后用。关于同余的介绍，请看模运算的世界：当加减乘除遇上取模。

本题有负数，根据这篇文章中的公式，我们可以把每个 $\textit{nums}[i]$ 变成

$$
(\textit{nums}[i]\bmod m + m)\bmod m
$$

从而保证取模结果在 $[0,m)$ 中。

例如 $\textit{nums}=[1,-6,-4,3,5]$，$m=3$，取模后变成 $[1,0,2,0,2]$。

然后枚举答案：

有没有与 $0$ 关于模 $m$ 同余的数？有，我们消耗掉一个 $0$。
有没有与 $1$ 关于模 $m$ 同余的数？有，我们消耗掉一个 $1$。
有没有与 $2$ 关于模 $m$ 同余的数？有，我们消耗掉一个 $2$。
有没有与 $3$ 关于模 $m$ 同余的数？有，我们消耗掉一个 $0$。这个取模后等于 $0$ 的数，可以继续操作，变成 $3$。
有没有与 $4$ 关于模 $m$ 同余的数？也就是看是否还有 $1$，没有，那么答案等于 $4$。

怎么知道还有没有剩余元素？用一个哈希表 $\textit{cnt}$ 统计 $(\textit{nums}[i]\bmod m + m) \bmod m$ 的个数。

本题视频讲解，欢迎点赞关注~

写法一

class Solution:
    def findSmallestInteger(self, nums: List[int], m: int) -> int:
        cnt = Counter(x % m for x in nums)
        mex = 0
        while cnt[mex % m]:
            cnt[mex % m] -= 1
            mex += 1
        return mex

class Solution {
    public int findSmallestInteger(int[] nums, int m) {
        int[] cnt = new int[m];
        for (int x : nums) {
            cnt[(x % m + m) % m]++; // 保证取模结果在 [0, m) 中
        }

        int mex = 0;
        while (cnt[mex % m]-- > 0) {
            mex++;
        }
        return mex;
    }
}

class Solution {
    public int findSmallestInteger(int[] nums, int m) {
        Map<Integer, Integer> cnt = new HashMap<>();
        for (int x : nums) {
            cnt.merge((x % m + m) % m, 1, Integer::sum);
        }

        int mex = 0;
        while (cnt.merge(mex % m, -1, Integer::sum) >= 0) {
            mex++;
        }
        return mex;
    }
}

class Solution {
public:
    int findSmallestInteger(vector<int>& nums, int m) {
        unordered_map<int, int> cnt;
        for (int x : nums) {
            cnt[(x % m + m) % m]++; // 保证取模结果在 [0, m) 中
        }

        int mex = 0;
        while (cnt[mex % m]-- > 0) {
            mex++;
        }
        return mex;
    }
};

int findSmallestInteger(int* nums, int numsSize, int m) {
    int* cnt = calloc(m, sizeof(int));
    for (int i = 0; i < numsSize; i++) {
        cnt[(nums[i] % m + m) % m]++; // 保证取模结果在 [0, m) 中
    }

    int mex = 0;
    while (cnt[mex % m]-- > 0) {
        mex++;
    }

    free(cnt);
    return mex;
}

func findSmallestInteger(nums []int, m int) (mex int) {
cnt := map[int]int{}
for _, x := range nums {
cnt[(x%m+m)%m]++ // 保证取模结果在 [0, m) 中
}

for cnt[mex%m] > 0 {
cnt[mex%m]--
mex++
}
return
}

var findSmallestInteger = function(nums, m) {
    const cnt = new Map();
    for (const x of nums) {
        const v = (x % m + m) % m; // 保证取模结果在 [0, m) 中
        cnt.set(v, (cnt.get(v) ?? 0) + 1);
    }

    let mex = 0;
    while ((cnt.get(mex % m) ?? 0) > 0) {
        cnt.set(mex % m, cnt.get(mex % m) - 1);
        mex++;
    }
    return mex;
};

use std::collections::HashMap;

impl Solution {
    pub fn find_smallest_integer(nums: Vec<i32>, m: i32) -> i32 {
        let mut cnt = HashMap::new();
        for x in nums {
            // 保证取模结果在 [0, m) 中
            *cnt.entry((x % m + m) % m).or_insert(0) += 1;
        }

        for mex in 0.. {
            if let Some(c) = cnt.get_mut(&(mex % m)) {
                if *c > 0 {
                    *c -= 1;
                    continue;
                }
            }
            return mex;
        }
        unreachable!()
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(n)$，其中 $n$ 是 $\textit{nums}$ 的长度。由于加多少个数，就只能减多少个数，所以第二个循环至多循环 $\mathcal{O}(n)$ 次。
空间复杂度：$\mathcal{O}(\min(n,m))$。哈希表中至多有 $\mathcal{O}(\min(n,m))$ 个元素。

写法二

{:width=500px}

此外，把哈希表换成数组更快。

class Solution:
    def findSmallestInteger(self, nums: List[int], m: int) -> int:
        cnt = [0] * m
        for x in nums:
            cnt[x % m] += 1

        i = cnt.index(min(cnt))
        return m * cnt[i] + i

class Solution {
    public int findSmallestInteger(int[] nums, int m) {
        int[] cnt = new int[m];
        for (int x : nums) {
            cnt[(x % m + m) % m]++;
        }

        int i = 0;
        for (int j = 1; j < m; j++) {
            if (cnt[j] < cnt[i]) {
                i = j;
            }
        }

        return m * cnt[i] + i;
    }
}

class Solution {
public:
    int findSmallestInteger(vector<int>& nums, int m) {
        vector<int> cnt(m);
        for (int x : nums) {
            cnt[(x % m + m) % m]++;
        }

        int i = ranges::min_element(cnt) - cnt.begin();
        return m * cnt[i] + i;
    }
};

int findSmallestInteger(int* nums, int numsSize, int m) {
    int* cnt = calloc(m, sizeof(int));
    for (int i = 0; i < numsSize; i++) {
        cnt[(nums[i] % m + m) % m]++;
    }

    int i = 0;
    for (int j = 1; j < m; j++) {
        if (cnt[j] < cnt[i]) {
            i = j;
        }
    }

    int ans = m * cnt[i] + i;
    free(cnt);
    return ans;
}

func findSmallestInteger(nums []int, m int) int {
cnt := make([]int, m)
for _, x := range nums {
cnt[(x%m+m)%m]++
}

i := 0
for j := 1; j < m; j++ {
if cnt[j] < cnt[i] {
i = j
}
}

return m*cnt[i] + i
}

var findSmallestInteger = function(nums, m) {
    const cnt = Array(m).fill(0);
    for (const x of nums) {
        cnt[(x % m + m) % m]++;
    }

    let i = 0;
    for (let j = 1; j < m; j++) {
        if (cnt[j] < cnt[i]) {
            i = j;
        }
    }

    return m * cnt[i] + i;
};

impl Solution {
    pub fn find_smallest_integer(nums: Vec<i32>, m: i32) -> i32 {
        let mut cnt = vec![0; m as usize];
        for x in nums {
            cnt[((x % m + m) % m) as usize] += 1;
        }

        let mut i = 0;
        for j in 1..m as usize {
            if cnt[j] < cnt[i] {
                i = j;
            }
        }

        m * cnt[i] + i as i32
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(n+m)$，其中 $n$ 是 $\textit{nums}$ 的长度。
空间复杂度：$\mathcal{O}(m)$。

相似题目

见下面数学题单的「§1.9 同余」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
O(n) 一次遍历，简洁写法（Python/Java/C++/C/Go/JS/Rust）endlesscheng
2024年11月10日 12:03

O(n) 一次遍历，简洁写法（Python/Java/C++/C/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2024年11月10日 12:03

遍历 $\textit{nums}$，寻找严格递增段（子数组）。

设当前严格递增段的长度为 $\textit{cnt}$，上一个严格递增段的长度为 $\textit{preCnt}$。

答案有两种情况：

两个子数组属于同一个严格递增段，那么 $k$ 最大是 $\left\lfloor\dfrac{\textit{cnt}}{2}\right\rfloor$。
两个子数组分别属于一对相邻的严格递增段，那么 $k$ 最大是 $\min(\textit{preCnt}, \textit{cnt})$。

本题视频讲解，欢迎点赞关注~

###py

class Solution:
    def maxIncreasingSubarrays(self, nums: List[int]) -> int:
        ans = pre_cnt = cnt = 0
        for i, x in enumerate(nums):
            cnt += 1
            if i == len(nums) - 1 or x >= nums[i + 1]:  # i 是严格递增段的末尾
                ans = max(ans, cnt // 2, min(pre_cnt, cnt))
                pre_cnt = cnt
                cnt = 0
        return ans

###java

class Solution {
    public int maxIncreasingSubarrays(List<Integer> nums) {
        int ans = 0;
        int preCnt = 0;
        int cnt = 0;
        for (int i = 0; i < nums.size(); i++) {
            cnt++;
            // i 是严格递增段的末尾
            if (i == nums.size() - 1 || nums.get(i) >= nums.get(i + 1)) {
                ans = Math.max(ans, Math.max(cnt / 2, Math.min(preCnt, cnt)));
                preCnt = cnt;
                cnt = 0;
            }
        }
        return ans;
    }
}

###java

class Solution {
    public int maxIncreasingSubarrays(List<Integer> nums) {
        Integer[] a = nums.toArray(Integer[]::new); // 转成数组处理，更快
        int ans = 0;
        int preCnt = 0;
        int cnt = 0;
        for (int i = 0; i < a.length; i++) {
            cnt++;
            // i 是严格递增段的末尾
            if (i == a.length - 1 || a[i] >= a[i + 1]) {
                ans = Math.max(ans, Math.max(cnt / 2, Math.min(preCnt, cnt)));
                preCnt = cnt;
                cnt = 0;
            }
        }
        return ans;
    }
}

###cpp

class Solution {
public:
    int maxIncreasingSubarrays(vector<int>& nums) {
        int ans = 0, pre_cnt = 0, cnt = 0;
        for (int i = 0; i < nums.size(); i++) {
            cnt++;
            if (i == nums.size() - 1 || nums[i] >= nums[i + 1]) { // i 是严格递增段的末尾
                ans = max({ans, cnt / 2, min(pre_cnt, cnt)});
                pre_cnt = cnt;
                cnt = 0;
            }
        }
        return ans;
    }
};

###c

#define MIN(a, b) ((b) < (a) ? (b) : (a))
#define MAX(a, b) ((b) > (a) ? (b) : (a))

int maxIncreasingSubarrays(int* nums, int numsSize) {
    int ans = 0, pre_cnt = 0, cnt = 0;
    for (int i = 0; i < numsSize; i++) {
        cnt++;
        if (i == numsSize - 1 || nums[i] >= nums[i + 1]) { // i 是严格递增段的末尾
            ans = MAX(ans, MAX(cnt / 2, MIN(pre_cnt, cnt)));
            pre_cnt = cnt;
            cnt = 0;
        }
    }
    return ans;
}

###go

func maxIncreasingSubarrays(nums []int) (ans int) {
preCnt, cnt := 0, 0
for i, x := range nums {
cnt++
if i == len(nums)-1 || x >= nums[i+1] { // i 是严格递增段的末尾
ans = max(ans, cnt/2, min(preCnt, cnt))
preCnt = cnt
cnt = 0
}
}
return
}

###js

var maxIncreasingSubarrays = function(nums) {
    let ans = 0, preCnt = 0, cnt = 0;
    for (let i = 0; i < nums.length; i++) {
        cnt++;
        if (i === nums.length - 1 || nums[i] >= nums[i + 1]) { // i 是严格递增段的末尾
            ans = Math.max(ans, Math.floor(cnt / 2), Math.min(preCnt, cnt));
            preCnt = cnt;
            cnt = 0;
        }
    }
    return ans;
};

###rust

impl Solution {
    pub fn max_increasing_subarrays(nums: Vec<i32>) -> i32 {
        let mut ans = 0;
        let mut pre_cnt = 0;
        let mut cnt = 0;
        for i in 0..nums.len() {
            cnt += 1;
            if i == nums.len() - 1 || nums[i] >= nums[i + 1] { // i 是严格递增段的末尾
                ans = ans.max(cnt / 2).max(pre_cnt.min(cnt));
                pre_cnt = cnt;
                cnt = 0;
            }
        }
        ans
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(n)$，其中 $n$ 是 $\textit{nums}$ 的长度。
空间复杂度：$\mathcal{O}(1)$。

专题训练

见下面双指针题单的「六、分组循环」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
O(n) 一次遍历（Python/Java/C++/Go）endlesscheng
2024年11月10日 12:04

O(n) 一次遍历（Python/Java/C++/Go）

LeetCode 每日一题题解

作者 endlesscheng

2024年11月10日 12:04

思路和周赛第二题是一样的，判断答案是否 $\ge k$，请看我的题解。

LeetCode 每日一题题解
从特殊到一般，教你如何思考本题，两种写法（Python/Java/C++/C/Go/JS/Rust）endlesscheng
2024年5月12日 12:44

从特殊到一般，教你如何思考本题，两种写法（Python/Java/C++/C/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2024年5月12日 12:44

k = 1

从特殊到一般，先想一想，$k=1$ 怎么做？

此时只能一步一步地向右走。无论起点在哪，终点都是 $n-1$。

如果选择 $i$ 为起点，我们计算的是子数组 $[i,n-1]$ 的元素和，即后缀和。

后缀和怎么算？我们可以倒着遍历 $\textit{energy}$，同时累加元素和，即为后缀和。

答案等于所有后缀和的最大值。

k = 2

再想一想，$k=2$ 怎么做？

此时我们有两个终点：$n-2$ 和 $n-1$。

对于终点 $n-1$：

如果选择 $n-3$ 为起点，那么我们累加的是下标为 $n-3,n-1$ 的元素和。
如果选择 $n-5$ 为起点，那么我们累加的是下标为 $n-5,n-3,n-1$ 的元素和。
如果选择 $n-7$ 为起点，那么我们累加的是下标为 $n-7,n-5,n-3,n-1$ 的元素和。
一般地，从 $n-1$ 开始倒着遍历，步长为 $-k=-2$，累加元素和，计算元素和的最大值。

对于终点 $n-2$：

如果选择 $n-4$ 为起点，那么我们累加的是下标为 $n-4,n-2$ 的元素和。
如果选择 $n-6$ 为起点，那么我们累加的是下标为 $n-6,n-4,n-2$ 的元素和。
如果选择 $n-8$ 为起点，那么我们累加的是下标为 $n-8,n-6,n-4,n-2$ 的元素和。
一般地，从 $n-2$ 开始倒着遍历，步长为 $-k=-2$，累加元素和，计算元素和的最大值。

是否可以从 $n-3$ 开始倒着遍历？

不行，因为 $n-3$ 还可以向右跳到 $n-1$，所以 $n-3$ 不是终点，不能作为倒着遍历的起点。

一般情况

枚举终点 $n-k,n-k+1,\dots,n-1$，倒着遍历，步长为 $-k$。

遍历的同时累加元素和 $\textit{sufSum}$，计算 $\textit{sufSum}$ 的最大值，即为答案。

写法一

###py

class Solution:
    def maximumEnergy(self, energy: List[int], k: int) -> int:
        n = len(energy)
        ans = -inf
        for i in range(n - k, n):  # 枚举终点 i
            suf_sum = accumulate(energy[j] for j in range(i, -1, -k))  # 计算后缀和
            ans = max(ans, max(suf_sum))
        return ans

###java

class Solution {
    public int maximumEnergy(int[] energy, int k) {
        int n = energy.length;
        int ans = Integer.MIN_VALUE;
        for (int i = n - k; i < n; i++) { // 枚举终点 i
            int sufSum = 0;
            for (int j = i; j >= 0; j -= k) {
                sufSum += energy[j]; // 计算后缀和
                ans = Math.max(ans, sufSum);
            }
        }
        return ans;
    }
}

###cpp

class Solution {
public:
    int maximumEnergy(vector<int>& energy, int k) {
        int n = energy.size();
        int ans = INT_MIN;
        for (int i = n - k; i < n; i++) { // 枚举终点 i
            int suf_sum = 0;
            for (int j = i; j >= 0; j -= k) {
                suf_sum += energy[j]; // 计算后缀和
                ans = max(ans, suf_sum);
            }
        }
        return ans;
    }
};

###c

#define MAX(a, b) ((b) > (a) ? (b) : (a))

int maximumEnergy(int* energy, int energySize, int k){
    int ans = INT_MIN;
    for (int i = energySize - k; i < energySize; i++) { // 枚举终点 i
        int suf_sum = 0;
        for (int j = i; j >= 0; j -= k) {
            suf_sum += energy[j]; // 计算后缀和
            ans = MAX(ans, suf_sum);
        }
    }
    return ans;
}

###go

func maximumEnergy(energy []int, k int) int {
n := len(energy)
ans := math.MinInt
for i := n - k; i < n; i++ { // 枚举终点 i
sufSum := 0
for j := i; j >= 0; j -= k {
sufSum += energy[j] // 计算后缀和
ans = max(ans, sufSum)
}
}
return ans
}

###js

var maximumEnergy = function(energy, k) {
    const n = energy.length;
    let ans = -Infinity;
    for (let i = n - k; i < n; i++) { // 枚举终点 i
        let sufSum = 0;
        for (let j = i; j >= 0; j -= k) {
            sufSum += energy[j]; // 计算后缀和
            ans = Math.max(ans, sufSum);
        }
    }
    return ans;
};

###rust

impl Solution {
    pub fn maximum_energy(energy: Vec<i32>, k: i32) -> i32 {
        let n = energy.len();
        let k = k as usize;
        let mut ans = i32::MIN;
        for i in n - k..n { // 枚举终点 i
            let mut suf_sum = 0;
            for j in (0..=i).rev().step_by(k) {
                suf_sum += energy[j]; // 计算后缀和
                ans = ans.max(suf_sum);
            }
        }
        ans
    }
}

写法二

原地计算后缀和，把后缀和保存到 $\textit{energy}$ 中。

最后返回 $\textit{energy}$ 的最大值，即为所有后缀和的最大值。

###py

class Solution:
    def maximumEnergy(self, energy: List[int], k: int) -> int:
        for i in range(len(energy) - k - 1, -1, -1):
            energy[i] += energy[i + k]
        return max(energy)

###java

class Solution {
    public int maximumEnergy(int[] energy, int k) {
        for (int i = energy.length - k - 1; i >= 0; i--) {
            energy[i] += energy[i + k];
        }
        return Arrays.stream(energy).max().getAsInt();
    }
}

###cpp

class Solution {
public:
    int maximumEnergy(vector<int>& energy, int k) {
        int n = energy.size();
        for (int i = n - k - 1; i >= 0; i--) {
            energy[i] += energy[i + k];
        }
        return ranges::max(energy);
    }
};

###c

#define MAX(a, b) ((b) > (a) ? (b) : (a))

int maximumEnergy(int* energy, int energySize, int k) {
    int ans = INT_MIN;
    for (int i = energySize - 1; i >= 0; i--) {
        if (i + k < energySize) {
            energy[i] += energy[i + k];
        }
        ans = MAX(ans, energy[i]);
    }
    return ans;
}

###go

func maximumEnergy(energy []int, k int) int {
for i := len(energy) - k - 1; i >= 0; i-- {
energy[i] += energy[i+k]
}
return slices.Max(energy)
}

###js

var maximumEnergy = function(energy, k) {
    for (let i = energy.length - k - 1; i >= 0; i--) {
        energy[i] += energy[i + k];
    }
    return Math.max(...energy);
};

###rust

impl Solution {
    pub fn maximum_energy(mut energy: Vec<i32>, k: i32) -> i32 {
        let k = k as usize;
        for i in (0..energy.len() - k).rev() {
            energy[i] += energy[i + k];
        }
        *energy.iter().max().unwrap()
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(n)$，其中 $n$ 是 $\textit{energy}$ 的长度。
空间复杂度：$\mathcal{O}(1)$。

专题训练

见下面贪心与思维题单的「§5.3 逆向思维」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
四种方法：正反两次扫描 / 递推 / record / 凸包+二分（Python/Java/C++/Go）endlesscheng
2025年3月23日 12:04

四种方法：正反两次扫描 / 递推 / record / 凸包+二分（Python/Java/C++/Go）

LeetCode 每日一题题解

作者 endlesscheng

2025年3月23日 12:04

方法一：正反两次扫描

为了计算酿造药水的时间，定义 $\textit{lastFinish}[i]$ 表示巫师 $i$ 完成上一瓶药水的时间。

示例 1 在处理完 $\textit{mana}[0]$ 后，有

$$
\textit{lastFinish} = [5,30,40,60]
$$

如果接着 $\textit{lastFinish}$ 继续酿造下一瓶药水 $\textit{mana}[1]=1$，完成时间是多少？注意开始酿造的时间不能早于 $\textit{lastFinish}[i]$。

$i$	$\textit{skill}[i]$	$\textit{lastFinish}[i]$	完成时间
$0$	$1$	$5$	$5+1=6$
$1$	$5$	$30$	$\max(6,30)+5=35$
$2$	$2$	$40$	$\max(35,40)+2=42$
$3$	$4$	$60$	$\max(42,60)+4=64$

题目要求「药水在当前巫师完成工作后必须立即传递给下一个巫师并开始处理」，也就是说，酿造药水的过程中是不能有停顿的。

从 $64$ 开始倒推，可以得到每名巫师的实际完成时间。比如倒数第二位巫师的完成时间，就是 $64$ 减去最后一名巫师花费的时间 $4\cdot 1$，得到 $60$。

$i$	$\textit{skill}[i+1]$	实际完成时间
$3$	-	$64$
$2$	$4$	$64-4\cdot 1=60$
$1$	$2$	$60-2\cdot 1=58$
$0$	$5$	$58-5\cdot 1=53$

按照上述过程处理每瓶药水，最终答案为 $\textit{lastFinish}[n-1]$。

本题视频讲解，欢迎点赞关注~

###py

class Solution:
    def minTime(self, skill: List[int], mana: List[int]) -> int:
        n = len(skill)
        last_finish = [0] * n  # 第 i 名巫师完成上一瓶药水的时间
        for m in mana:
            # 按题意模拟
            sum_t = 0
            for x, last in zip(skill, last_finish):
                if last > sum_t: sum_t = last  # 手写 max
                sum_t += x * m
            # 倒推：如果酿造药水的过程中没有停顿，那么 last_finish[i] 应该是多少
            last_finish[-1] = sum_t
            for i in range(n - 2, -1, -1):
                last_finish[i] = last_finish[i + 1] - skill[i + 1] * m
        return last_finish[-1]

###java

class Solution {
    public long minTime(int[] skill, int[] mana) {
        int n = skill.length;
        long[] lastFinish = new long[n]; // 第 i 名巫师完成上一瓶药水的时间
        for (int m : mana) {
            // 按题意模拟
            long sumT = 0;
            for (int i = 0; i < n; i++) {
                sumT = Math.max(sumT, lastFinish[i]) + skill[i] * m;
            }
            // 倒推：如果酿造药水的过程中没有停顿，那么 lastFinish[i] 应该是多少
            lastFinish[n - 1] = sumT;
            for (int i = n - 2; i >= 0; i--) {
                lastFinish[i] = lastFinish[i + 1] - skill[i + 1] * m;
            }
        }
        return lastFinish[n - 1];
    }
}

###cpp

class Solution {
public:
    long long minTime(vector<int>& skill, vector<int>& mana) {
        int n = skill.size();
        vector<long long> last_finish(n); // 第 i 名巫师完成上一瓶药水的时间
        for (int m : mana) {
            // 按题意模拟
            long long sum_t = 0;
            for (int i = 0; i < n; i++) {
                sum_t = max(sum_t, last_finish[i]) + skill[i] * m;
            }
            // 倒推：如果酿造药水的过程中没有停顿，那么 lastFinish[i] 应该是多少
            last_finish[n - 1] = sum_t;
            for (int i = n - 2; i >= 0; i--) {
                last_finish[i] = last_finish[i + 1] - skill[i + 1] * m;
            }
        }
        return last_finish[n - 1];
    }
};

###go

func minTime(skill, mana []int) int64 {
n := len(skill)
lastFinish := make([]int, n) // 第 i 名巫师完成上一瓶药水的时间
for _, m := range mana {
// 按题意模拟
sumT := 0
for i, x := range skill {
sumT = max(sumT, lastFinish[i]) + x*m
}
// 倒推：如果酿造药水的过程中没有停顿，那么 lastFinish[i] 应该是多少
lastFinish[n-1] = sumT
for i := n - 2; i >= 0; i-- {
lastFinish[i] = lastFinish[i+1] - skill[i+1]*m
}
}
return int64(lastFinish[n-1])
}

复杂度分析

时间复杂度：$\mathcal{O}(nm)$，其中 $n$ 是 $\textit{skill}$ 的长度，$m$ 是 $\textit{mana}$ 的长度。
空间复杂度：$\mathcal{O}(n)$。

方法二：递推开始时间

由于酿造药水的过程是连续的，所以知道了开始时间（或者完成时间）就能知道每个 $\textit{lastFinish}[i]$。所以 $\textit{lastFinish}$ 数组是多余的。

设开始酿造 $\textit{mana}[j]$ 的时间为 $\textit{start}_j$，那么有

$$
\textit{lastFinish}_j[i] = \textit{start}j + \textit{mana}[j]\cdot \sum{k=0}^{i} \textit{skill}[k]
$$

在已知 $\textit{start}_{j-1}$ 的前提下，能否递推算出 $\textit{start}_j$？

哪位巫师决定了开始时间？假设第 $i$ 位巫师决定了开始时间，那么这位巫师完成 $\textit{mana}[j-1]$ 的时间，同时也是他开始 $\textit{mana}[j]$ 的时间。

所以有

$$
\textit{lastFinish}_{j-1}[i] + \textit{mana}[j]\cdot \textit{skill}[i] = \textit{lastFinish}_j[i]
$$

两边代入 $\textit{lastFinish}_j[i]$ 的式子，得

$$
\textit{start}{j-1} + \textit{mana}[j-1]\cdot \sum{k=0}^{i} \textit{skill}[k] + \textit{mana}[j]\cdot \textit{skill}[i] = \textit{start}j + \textit{mana}[j]\cdot \sum{k=0}^{i} \textit{skill}[k]
$$

移项得

$$
\textit{start}j = \textit{start}{j-1} + \textit{mana}[j-1]\cdot \sum_{k=0}^{i} \textit{skill}[k] - \textit{mana}[j]\cdot \sum_{k=0}^{i-1} \textit{skill}[k]
$$

计算 $\textit{skill}$ 的前缀和数组 $s$，上式为

$$
\textit{start}j = \textit{start}{j-1} + \textit{mana}[j-1]\cdot s[i+1] - \textit{mana}[j]\cdot s[i]
$$

枚举 $i$，取最大值，得

$$
\textit{start}j = \textit{start}{j-1} + \max_{i=0}^{n-1} \left{\textit{mana}[j-1]\cdot s[i+1] - \textit{mana}[j]\cdot s[i]\right}
$$

初始值 $\textit{start}_0 = 0$。

答案为 $\textit{lastFinish}{m-1}[n-1] = \textit{start}{m-1} + \textit{mana}[m-1]\cdot s[n]$。

###py

class Solution:
    def minTime(self, skill: List[int], mana: List[int]) -> int:
        n = len(skill)
        s = list(accumulate(skill, initial=0))  # skill 的前缀和
        start = 0
        for pre, cur in pairwise(mana):
            start += max(pre * s[i + 1] - cur * s[i] for i in range(n))
        return start + mana[-1] * s[-1]

###java

class Solution {
    public long minTime(int[] skill, int[] mana) {
        int n = skill.length;
        int[] s = new int[n + 1]; // skill 的前缀和
        for (int i = 0; i < n; i++) {
            s[i + 1] = s[i] + skill[i];
        }

        int m = mana.length;
        long start = 0;
        for (int j = 1; j < m; j++) {
            long mx = 0;
            for (int i = 0; i < n; i++) {
                mx = Math.max(mx, (long) mana[j - 1] * s[i + 1] - (long) mana[j] * s[i]);
            }
            start += mx;
        }
        return start + (long) mana[m - 1] * s[n];
    }
}

###cpp

class Solution {
public:
    long long minTime(vector<int>& skill, vector<int>& mana) {
        int n = skill.size(), m = mana.size();
        vector<int> s(n + 1); // skill 的前缀和
        partial_sum(skill.begin(), skill.end(), s.begin() + 1);

        long long start = 0;
        for (int j = 1; j < m; j++) {
            long long mx = 0;
            for (int i = 0; i < n; i++) {
                mx = max(mx, 1LL * mana[j - 1] * s[i + 1] - 1LL * mana[j] * s[i]);
            }
            start += mx;
        }
        return start + 1LL * mana[m - 1] * s[n];
    }
};

###go

func minTime(skill, mana []int) int64 {
n, m := len(skill), len(mana)
s := make([]int, n+1) // skill 的前缀和
for i, x := range skill {
s[i+1] = s[i] + x
}

start := 0
for j := 1; j < m; j++ {
mx := 0
for i := range n {
mx = max(mx, mana[j-1]*s[i+1]-mana[j]*s[i])
}
start += mx
}
return int64(start + mana[m-1]*s[n])
}

复杂度分析

时间复杂度：$\mathcal{O}(nm)$，其中 $n$ 是 $\textit{skill}$ 的长度，$m$ 是 $\textit{mana}$ 的长度。
空间复杂度：$\mathcal{O}(n)$。如果在遍历的同时计算前缀和，则可以做到 $\mathcal{O}(1)$ 空间。

方法三：record 优化

将递推式

$$
\textit{start}j = \textit{start}{j-1} + \max_{i=0}^{n-1} \left{\textit{mana}[j-1]\cdot s[i+1] - \textit{mana}[j]\cdot s[i]\right}
$$

变形为

$$
\textit{start}j = \textit{start}{j-1} + \max_{i=0}^{n-1} \left{(\textit{mana}[j-1]-\textit{mana}[j])\cdot s[i] + \textit{mana}[j-1]\cdot \textit{skill}[i] \right}
$$

设 $d = \textit{mana}[j-1]-\textit{mana}[j]$。分类讨论：

如果 $d > 0$。由于 $s$ 是单调递增数组，如果 $\textit{skill}[3] < \textit{skill}[5]$，那么 $i=3$ 绝对不会算出最大值；但如果 $\textit{skill}[3] > \textit{skill}[5]$，谁会算出最大值就不一定了。所以我们只需要考虑 $\textit{skill}$ 的逆序 record，这才是可能成为最大值的数据。其中逆序 record 的意思是，倒序遍历 $\textit{skill}$，每次遍历到更大的数，就记录下标。
如果 $d < 0$。由于 $s$ 是单调递增数组，如果 $\textit{skill}[5] < \textit{skill}[3]$，那么 $i=5$ 绝对不会算出最大值；但如果 $\textit{skill}[5] > \textit{skill}[3]$，谁会算出最大值就不一定了。所以我们只需要考虑 $\textit{skill}$ 的正序 record，这才是可能成为最大值的数据。其中正序 record 的意思是，正序遍历 $\textit{skill}$，每次遍历到更大的数，就记录下标。
$d = 0$ 的情况可以并入 $d>0$ 的情况。

###py

class Solution:
    def minTime(self, skill: List[int], mana: List[int]) -> int:
        n = len(skill)
        s = list(accumulate(skill, initial=0))

        suf_record = [n - 1]
        for i in range(n - 2, -1, -1):
            if skill[i] > skill[suf_record[-1]]:
                suf_record.append(i)

        pre_record = [0]
        for i in range(1, n):
            if skill[i] > skill[pre_record[-1]]:
                pre_record.append(i)

        start = 0
        for pre, cur in pairwise(mana):
            record = pre_record if pre < cur else suf_record
            start += max(pre * s[i + 1] - cur * s[i] for i in record)
        return start + mana[-1] * s[-1]

###java

class Solution {
    public long minTime(int[] skill, int[] mana) {
        int n = skill.length;
        int[] s = new int[n + 1];
        for (int i = 0; i < n; i++) {
            s[i + 1] = s[i] + skill[i];
        }

        List<Integer> suf = new ArrayList<>();
        suf.add(n - 1);
        for (int i = n - 2; i >= 0; i--) {
            if (skill[i] > skill[suf.getLast()]) {
                suf.add(i);
            }
        }

        List<Integer> pre = new ArrayList<>();
        pre.add(0);
        for (int i = 1; i < n; i++) {
            if (skill[i] > skill[pre.getLast()]) {
                pre.add(i);
            }
        }

        int m = mana.length;
        long start = 0;
        for (int j = 1; j < m; j++) {
            List<Integer> record = mana[j - 1] < mana[j] ? pre : suf;
            long mx = 0;
            for (int i : record) {
                mx = Math.max(mx, (long) mana[j - 1] * s[i + 1] - (long) mana[j] * s[i]);
            }
            start += mx;
        }
        return start + (long) mana[m - 1] * s[n];
    }
}

###java

class Solution {
    public long minTime(int[] skill, int[] mana) {
        int n = skill.length;
        int[] s = new int[n + 1];
        for (int i = 0; i < n; i++) {
            s[i + 1] = s[i] + skill[i];
        }

        int[] suf = new int[n];
        int sufLen = 0;
        suf[sufLen++] = n - 1;
        for (int i = n - 2; i >= 0; i--) {
            if (skill[i] > skill[suf[sufLen - 1]]) {
                suf[sufLen++] = i;
            }
        }

        int[] pre = new int[n];
        int preLen = 0;
        pre[preLen++] = 0;
        for (int i = 1; i < n; i++) {
            if (skill[i] > skill[pre[preLen - 1]]) {
                pre[preLen++] = i;
            }
        }

        int m = mana.length;
        long start = 0;
        for (int j = 1; j < m; j++) {
            int[] record = mana[j - 1] < mana[j] ? pre : suf;
            int recordLen = mana[j - 1] < mana[j] ? preLen : sufLen;
            long mx = 0;
            for (int k = 0; k < recordLen; k++) {
                int i = record[k];
                mx = Math.max(mx, (long) mana[j - 1] * s[i + 1] - (long) mana[j] * s[i]);
            }
            start += mx;
        }
        return start + (long) mana[m - 1] * s[n];
    }
}

###cpp

class Solution {
public:
    long long minTime(vector<int>& skill, vector<int>& mana) {
        int n = skill.size(), m = mana.size();
        vector<int> s(n + 1);
        partial_sum(skill.begin(), skill.end(), s.begin() + 1);

        vector<int> suf = {n - 1};
        for (int i = n - 2; i >= 0; i--) {
            if (skill[i] > skill[suf.back()]) {
                suf.push_back(i);
            }
        }

        vector<int> pre = {0};
        for (int i = 1; i < n; i++) {
            if (skill[i] > skill[pre.back()]) {
                pre.push_back(i);
            }
        }

        long long start = 0;
        for (int j = 1; j < m; j++) {
            auto& record = mana[j - 1] < mana[j] ? pre : suf;
            long long mx = 0;
            for (int i : record) {
                mx = max(mx, 1LL * mana[j - 1] * s[i + 1] - 1LL * mana[j] * s[i]);
            }
            start += mx;
        }
        return start + 1LL * mana[m - 1] * s[n];
    }
};

###go

func minTime(skill, mana []int) int64 {
n, m := len(skill), len(mana)
s := make([]int, n+1)
for i, x := range skill {
s[i+1] = s[i] + x
}

suf := []int{n - 1}
for i := n - 2; i >= 0; i-- {
if skill[i] > skill[suf[len(suf)-1]] {
suf = append(suf, i)
}
}

pre := []int{0}
for i := 1; i < n; i++ {
if skill[i] > skill[pre[len(pre)-1]] {
pre = append(pre, i)
}
}

start := 0
for j := 1; j < m; j++ {
record := suf
if mana[j-1] < mana[j] {
record = pre
}
mx := 0
for _, i := range record {
mx = max(mx, mana[j-1]*s[i+1]-mana[j]*s[i])
}
start += mx
}
return int64(start + mana[m-1]*s[n])
}

复杂度分析（最坏情况）

时间复杂度：$\mathcal{O}(nm)$，其中 $n$ 是 $\textit{skill}$ 的长度，$m$ 是 $\textit{mana}$ 的长度。
空间复杂度：$\mathcal{O}(n)$。

复杂度分析（平均情况）

力扣喜欢出随机数据，上述算法在随机数据下的性能如何？

换句话说，在随机数据下，record 的期望长度是多少？

为方便分析，假设 $\textit{skill}$ 是一个随机的 $[1,n]$ 的排列。

$\textit{skill}[i]$ 如果是一个新的最大值，那么它是 $[0,i]$ 中的最大值。在随机排列的情况下，$[0,i]$ 中的排列也是随机的，所以这等价于该排列的最后一个数是最大值的概率，即

$$
\dfrac{1}{i+1}
$$

record 的期望长度，等于「每个位置能否成为新的最大值」之和，能就贡献 $1$，不能就贡献 $0$。

所以 $\textit{skill}[i]$ 给期望的贡献是 $\dfrac{1}{i+1}$。所以 record 的期望长度为

$$
\sum_{i=0}^{n-1} \dfrac{1}{i+1}
$$

由调和级数可知，record 的期望长度为 $\Theta(\log n)$。

时间复杂度：$\Theta(n + m\log n)$，其中 $n$ 是 $\textit{skill}$ 的长度，$m$ 是 $\textit{mana}$ 的长度。
空间复杂度：$\mathcal{O}(n)$。

方法四：凸包 + 二分

前置知识：二维计算几何，凸包，Andrew 算法。

把递推式

$$
\textit{start}j = \textit{start}{j-1} + \max_{i=0}^{n-1} \left{(\textit{mana}[j-1]-\textit{mana}[j])\cdot s[i] + \textit{mana}[j-1]\cdot \textit{skill}[i] \right}
$$

中的

$$
(\textit{mana}[j-1]-\textit{mana}[j])\cdot s[i] + \textit{mana}[j-1]\cdot \textit{skill}[i]
$$

改成点积的形式，这样我们能得到来自几何意义上的观察。

设向量 $\mathbf{v}_i = (s[i],\textit{skill}[i])$。

设向量 $\mathbf{p} = (\textit{mana}[j-1]-\textit{mana}[j], \textit{mana}[j-1])$。

那么我们求的是

$$
\max_{i=0}^{n-1} \mathbf{p}\cdot \mathbf{v}_i
$$

根据点积的几何意义，我们求的是 $\mathbf{v}_i$ 在 $\mathbf{p}$ 方向上的投影长度，再乘以 $\mathbf{p}$ 的模长 $||\mathbf{p}||$。由于 $||\mathbf{p}||$ 是个定值，所以要最大化投影长度。

考虑 $\mathbf{v}_i$ 的上凸包（用 Andrew 算法计算），在凸包内的点，就像是山坳，比凸包顶点的投影长度短。所以只需考虑凸包顶点。

这样有一个很好的性质：顺时针（或者逆时针）遍历凸包顶点，$\mathbf{p}\cdot \mathbf{v}_i$ 会先变大再变小（单峰函数）。那么要计算最大值，就类似 852. 山脉数组的峰顶索引，二分首个「下坡」的位置，具体见我的题解。

###py

class Vec:
    __slots__ = 'x', 'y'

    def __init__(self, x: int, y: int):
        self.x = x
        self.y = y

    def __sub__(self, b: "Vec") -> "Vec":
        return Vec(self.x - b.x, self.y - b.y)

    def det(self, b: "Vec") -> int:
        return self.x * b.y - self.y * b.x

    def dot(self, b: "Vec") -> int:
        return self.x * b.x + self.y * b.y

class Solution:
    # Andrew 算法，计算 points 的上凸包
    # 由于横坐标（前缀和）是严格递增的，所以无需排序
    def convex_hull(self, points: List[Vec]) -> List[Vec]:
        q = []
        for p in points:
            while len(q) > 1 and (q[-1] - q[-2]).det(p - q[-1]) >= 0:
                q.pop()
            q.append(p)
        return q

    def minTime(self, skill: List[int], mana: List[int]) -> int:
        s = list(accumulate(skill, initial=0))
        vs = [Vec(pre_sum, x) for pre_sum, x in zip(s, skill)]
        vs = self.convex_hull(vs)  # 去掉无用数据

        start = 0
        for pre, cur in pairwise(mana):
            p = Vec(pre - cur, pre)
            # p.dot(vs[i]) 是个单峰函数，二分找最大值
            check = lambda i: p.dot(vs[i]) > p.dot(vs[i + 1])
            i = bisect_left(range(len(vs) - 1), True, key=check)
            start += p.dot(vs[i])
        return start + mana[-1] * s[-1]

###java

class Solution {
    private record Vec(int x, int y) {
        Vec sub(Vec b) {
            return new Vec(x - b.x, y - b.y);
        }

        long det(Vec b) {
            return (long) x * b.y - (long) y * b.x;
        }

        long dot(Vec b) {
            return (long) x * b.x + (long) y * b.y;
        }
    }

    // Andrew 算法，计算 points 的上凸包
    // 由于横坐标（前缀和）是严格递增的，所以无需排序
    private List<Vec> convexHull(Vec[] points) {
        List<Vec> q = new ArrayList<>();
        for (Vec p : points) {
            while (q.size() > 1 && q.getLast().sub(q.get(q.size() - 2)).det(p.sub(q.getLast())) >= 0) {
                q.removeLast();
            }
            q.add(p);
        }
        return q;
    }

    public long minTime(int[] skill, int[] mana) {
        int n = skill.length;
        int[] s = new int[n + 1];
        Vec[] points = new Vec[n];
        for (int i = 0; i < n; i++) {
            s[i + 1] = s[i] + skill[i];
            points[i] = new Vec(s[i], skill[i]);
        }
        List<Vec> vs = convexHull(points); // 去掉无用数据

        int m = mana.length;
        long start = 0;
        for (int j = 1; j < m; j++) {
            Vec p = new Vec(mana[j - 1] - mana[j], mana[j - 1]);
            // p.dot(vs[i]) 是个单峰函数，二分找最大值
            int l = -1, r = vs.size() - 1;
            while (l + 1 < r) {
                int mid = (l + r) >>> 1;
                if (p.dot(vs.get(mid)) > p.dot(vs.get(mid + 1))) {
                    r = mid;
                } else {
                    l = mid;
                }
            }
            start += p.dot(vs.get(r));
        }
        return start + (long) mana[m - 1] * s[n];
    }
}

###cpp

struct Vec {
    int x, y;
    Vec operator-(const Vec& b) { return {x - b.x, y - b.y}; }
    long long det(const Vec& b) { return 1LL * x * b.y - 1LL * y * b.x; }
    long long dot(const Vec& b) { return 1LL * x * b.x + 1LL * y * b.y; }
};

class Solution {
    // Andrew 算法，计算 points 的上凸包
    // 由于横坐标（前缀和）是严格递增的，所以无需排序
    vector<Vec> convex_hull(vector<Vec>& points) {
        vector<Vec> q;
        for (auto& p : points) {
            while (q.size() > 1 && (q.back() - q[q.size() - 2]).det(p - q.back()) >= 0) {
                q.pop_back();
            }
            q.push_back(p);
        }
        return q;
    }

public:
    long long minTime(vector<int>& skill, vector<int>& mana) {
        int n = skill.size(), m = mana.size();
        vector<int> s(n + 1);
        vector<Vec> vs(n);
        for (int i = 0; i < n; i++) {
            s[i + 1] = s[i] + skill[i];
            vs[i] = {s[i], skill[i]};
        }
        vs = convex_hull(vs); // 去掉无用数据

        long long start = 0;
        for (int j = 1; j < m; j++) {
            Vec p = {mana[j - 1] - mana[j], mana[j - 1]};
            // p.dot(vs[i]) 是个单峰函数，二分找最大值
            int l = -1, r = vs.size() - 1;
            while (l + 1 < r) {
                int mid = l + (r - l) / 2;
                (p.dot(vs[mid]) > p.dot(vs[mid + 1]) ? r : l) = mid;
            }
            start += p.dot(vs[r]);
        }
        return start + 1LL * mana[m - 1] * s[n];
    }
};

###go

type vec struct{ x, y int }

func (a vec) sub(b vec) vec { return vec{a.x - b.x, a.y - b.y} }
func (a vec) det(b vec) int { return a.x*b.y - a.y*b.x }
func (a vec) dot(b vec) int { return a.x*b.x + a.y*b.y }

// Andrew 算法，计算 points 的上凸包
// 由于横坐标（前缀和）是严格递增的，所以无需排序
func convexHull(points []vec) (q []vec) {
for _, p := range points {
for len(q) > 1 && q[len(q)-1].sub(q[len(q)-2]).det(p.sub(q[len(q)-1])) >= 0 {
q = q[:len(q)-1]
}
q = append(q, p)
}
return
}

func minTime(skill, mana []int) int64 {
n, m := len(skill), len(mana)
s := make([]int, n+1)
vs := make([]vec, n)
for i, x := range skill {
s[i+1] = s[i] + x
vs[i] = vec{s[i], x}
}
vs = convexHull(vs) // 去掉无用数据

start := 0
for j := 1; j < m; j++ {
p := vec{mana[j-1] - mana[j], mana[j-1]}
// p.dot(vs[i]) 是个单峰函数，二分找最大值
i := sort.Search(len(vs)-1, func(i int) bool { return p.dot(vs[i]) > p.dot(vs[i+1]) })
start += p.dot(vs[i])
}
return int64(start + mana[m-1]*s[n])
}

复杂度分析

时间复杂度：$\mathcal{O}(n + m\log n)$，其中 $n$ 是 $\textit{skill}$ 的长度，$m$ 是 $\textit{mana}$ 的长度。
空间复杂度：$\mathcal{O}(n)$。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
两种方法：排序+二分 / 计数+值域后缀和（Python/Java/C++/C/Go/JS/Rust）endlesscheng
2022年6月12日 07:00

两种方法：排序+二分 / 计数+值域后缀和（Python/Java/C++/C/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2022年6月12日 07:00

方法一：排序 + 二分查找

写法一：使用浮点数

问题相当于给你 $n$ 个询问，每次问 $\textit{spells}[i]$ 与 $\textit{potions}$ 中的多少个数相乘，结果 $\ge \textit{success}$。

对 $\textit{potions}$ 排序后，就可以二分查找了：

设 $j$ 是最小的满足 $\textit{potions}[j]\ge\dfrac{\textit{success}}{\textit{spells}[i]}$ 的下标。
由于数组已经排序，那么下标大于 $j$ 的数也同样满足不等式。
从 $j$ 到 $m-1$，一共有 $m-j$ 个满足不等式的数，其中 $m$ 是 $\textit{potions}$ 的长度。

有关二分查找的原理，请看【基础算法精讲 04】。

class Solution:
    def successfulPairs(self, spells: List[int], potions: List[int], success: int) -> List[int]:
        potions.sort()
        m = len(potions)
        return [m - bisect_left(potions, success / x) for x in spells]

class Solution {
    public int[] successfulPairs(int[] spells, int[] potions, long success) {
        Arrays.sort(potions);
        for (int i = 0; i < spells.length; i++) {
            spells[i] = potions.length - lowerBound(potions, (double) success / spells[i]);
        }
        return spells;
    }

    // 返回 nums 中的第一个大于等于 target 的元素下标
    private int lowerBound(int[] nums, double target) {
        int left = -1, right = nums.length; // 开区间 (left, right)
        while (left + 1 < right) { // 区间不为空
            // 循环不变量：
            // nums[left] <= target
            // nums[right] > target
            int mid = (left + right) >>> 1; // 比 left+(right-left)/2 更快的写法
            if (nums[mid] >= target) {
                right = mid; // 二分范围缩小到 (left, mid)
            } else {
                left = mid; // 二分范围缩小到 (mid, right)
            }
        }
        return right;
    }
}

class Solution {
public:
    vector<int> successfulPairs(vector<int>& spells, vector<int>& potions, long long success) {
        ranges::sort(potions);
        for (int& x : spells) { // 原地修改
            x = potions.end() - ranges::lower_bound(potions, 1.0 * success / x);
        }
        return spells;
    }
};

int cmp(const void* a, const void* b) {
    return *(int*)a - *(int*)b;
}

// 返回 nums 中的第一个大于等于 target 的元素下标
int lowerBound(int* nums, int numsSize, double target) {
    int left = -1, right = numsSize; // 开区间 (left, right)
    while (left + 1 < right) { // 区间不为空
        int mid = left + (right - left) / 2;
        if (nums[mid] >= target) {
            right = mid; // 二分范围缩小到 (left, mid)
        } else {
            left = mid; // 二分范围缩小到 (mid, right)
        }
    }
    return right;
}

int* successfulPairs(int* spells, int spellsSize, int* potions, int potionsSize, long long success, int* returnSize) {
    qsort(potions, potionsSize, sizeof(int), cmp);
    for (int i = 0; i < spellsSize; i++) {
        spells[i] = potionsSize - lowerBound(potions, potionsSize, 1.0 * success / spells[i]);
    }
    *returnSize = spellsSize;
    return spells;
}

func successfulPairs(spells, potions []int, success int64) []int {
    slices.Sort(potions)
    for i, x := range spells {
        target := float64(success) / float64(x)
        j := sort.Search(len(potions), func(j int) bool { return float64(potions[j]) >= target })
        spells[i] = len(potions) - j
    }
    return spells
}

var successfulPairs = function(spells, potions, success) {
    potions.sort((a, b) => a - b);
    for (let i = 0; i < spells.length; i++) {
        const target = success / spells[i];
        spells[i] = potions.length - lowerBound(potions, target);
    }
    return spells;
};

var lowerBound = function(nums, target) {
    let left = -1, right = nums.length; // 开区间 (left, right)
    while (left + 1 < right) { // 区间不为空
        // 循环不变量：
        // nums[left] < target
        // nums[right] >= target
        const mid = Math.floor((left + right) / 2);
        if (nums[mid] >= target) {
            right = mid; // 范围缩小到 (left, mid)
        } else {
            left = mid; // 范围缩小到 (mid, right)
        }
    }
    return right;
}

var successfulPairs = function(spells, potions, success) {
    potions.sort((a, b) => a - b);
    for (let i = 0; i < spells.length; i++) {
        const target = success / spells[i];
        spells[i] = potions.length - _.sortedIndex(potions, target);
    }
    return spells;
};

impl Solution {
    pub fn successful_pairs(mut spells: Vec<i32>, mut potions: Vec<i32>, success: i64) -> Vec<i32> {
        potions.sort_unstable();
        let last = potions[potions.len() - 1] as i64;
        for x in spells.iter_mut() {
            let target = success as f64 / *x as f64;
            let j = potions.partition_point(|&x| (x as f64) < target);
            *x = (potions.len() - j) as i32;
        }
        spells
    }
}

写法二：不使用浮点数

浮点数有舍入误差，如果数据范围更大，上面的做法就不一定正确了。

更好的做法是，避免使用浮点数，只使用整数计算。一方面可以保证正确性，另一方面整数运算比浮点运算快。

对于正整数，$xy\ge\textit{success}$ 等价于 $y\ge\left\lceil\dfrac{\textit{success}}{x}\right\rceil$。

为方便二分，可以利用如下恒等式：

$$
\left\lceil\dfrac{a}{b}\right\rceil = \left\lfloor\dfrac{a+b-1}{b}\right\rfloor = \left\lfloor\dfrac{a-1}{b}\right\rfloor + 1
$$

证明见上取整下取整转换公式的证明。

根据上式，我们有

$$
y\ge\left\lceil\dfrac{\textit{success}}{x}\right\rceil = \left\lfloor\dfrac{\textit{success}-1}{x}\right\rfloor + 1
$$

也可以写成

$$
y>\left\lfloor\dfrac{\textit{success}-1}{x}\right\rfloor
$$

为什么不等式一定要这样变形？好处是只需要在二分之前做一次除法，避免在二分循环内计算乘法，效率更高。另外的好处是部分语言可以直接调用库函数二分。

class Solution:
    def successfulPairs(self, spells: List[int], potions: List[int], success: int) -> List[int]:
        potions.sort()
        m = len(potions)
        success -= 1  # 提前减一，避免在循环中反复减一
        return [m - bisect_right(potions, success // x) for x in spells]

class Solution {
    public int[] successfulPairs(int[] spells, int[] potions, long success) {
        Arrays.sort(potions);
        for (int i = 0; i < spells.length; i++) {
            long target = (success - 1) / spells[i];
            if (target < potions[potions.length - 1]) {
                // 这样写每次二分就只用比两个 int 的大小，避免把 potions 中的元素转成 long 比较
                spells[i] = potions.length - upperBound(potions, (int) target);
            } else {
                spells[i] = 0;
            }
        }
        return spells;
    }

    // 返回 nums 中的第一个大于 target 的元素下标
    private int upperBound(int[] nums, int target) {
        int left = -1, right = nums.length; // 开区间 (left, right)
        while (left + 1 < right) { // 区间不为空
            // 循环不变量：
            // nums[left] <= target
            // nums[right] > target
            int mid = (left + right) >>> 1; // 比 left+(right-left)/2 更快的写法
            if (nums[mid] > target) {
                right = mid; // 二分范围缩小到 (left, mid)
            } else {
                left = mid; // 二分范围缩小到 (mid, right)
            }
        }
        return right;
    }
}

class Solution {
public:
    vector<int> successfulPairs(vector<int>& spells, vector<int>& potions, long long success) {
        ranges::sort(potions);
        for (int& x : spells) { // 原地修改
            long long target = (success - 1) / x;
            if (target < potions.back()) {
                // 这样写每次二分就只用比两个 int 的大小，避免把 potions 中的元素转成 long long 比较
                x = potions.end() - ranges::upper_bound(potions, (int) target);
            } else {
                x = 0;
            }
        }
        return spells;
    }
};

int cmp(const void* a, const void* b) {
    return *(int*)a - *(int*)b;
}

// 返回 nums 中的第一个大于 target 的元素下标
int upperBound(int* nums, int numsSize, int target) {
    int left = -1, right = numsSize; // 开区间 (left, right)
    while (left + 1 < right) { // 区间不为空
        int mid = left + (right - left) / 2;
        if (nums[mid] > target) {
            right = mid; // 二分范围缩小到 (left, mid)
        } else {
            left = mid; // 二分范围缩小到 (mid, right)
        }
    }
    return right;
}

int* successfulPairs(int* spells, int spellsSize, int* potions, int potionsSize, long long success, int* returnSize) {
    qsort(potions, potionsSize, sizeof(int), cmp);
    for (int i = 0; i < spellsSize; i++) {
        long long target = (success - 1) / spells[i];
        if (target < potions[potionsSize - 1]) {
            // 这样写每次二分就只用比两个 int 的大小，避免把 potions 中的元素转成 long long 比较
            spells[i] = potionsSize - upperBound(potions, potionsSize, target);
        } else {
            spells[i] = 0;
        }
    }
    *returnSize = spellsSize;
    return spells;
}

func successfulPairs(spells, potions []int, success int64) []int {
slices.Sort(potions)
for i, x := range spells {
spells[i] = len(potions) - sort.SearchInts(potions, (int(success)-1)/x+1)
}
return spells
}

var successfulPairs = function(spells, potions, success) {
    potions.sort((a, b) => a - b);
    for (let i = 0; i < spells.length; i++) {
        const target = Math.ceil(success / spells[i]);
        spells[i] = potions.length - lowerBound(potions, target);
    }
    return spells;
};

var lowerBound = function(nums, target) {
    let left = -1, right = nums.length; // 开区间 (left, right)
    while (left + 1 < right) { // 区间不为空
        // 循环不变量：
        // nums[left] < target
        // nums[right] >= target
        const mid = Math.floor((left + right) / 2);
        if (nums[mid] >= target) {
            right = mid; // 范围缩小到 (left, mid)
        } else {
            left = mid; // 范围缩小到 (mid, right)
        }
    }
    return right;
}

var successfulPairs = function(spells, potions, success) {
    potions.sort((a, b) => a - b);
    for (let i = 0; i < spells.length; i++) {
        const target = Math.ceil(success / spells[i]);
        spells[i] = potions.length - _.sortedIndex(potions, target);
    }
    return spells;
};

impl Solution {
    pub fn successful_pairs(mut spells: Vec<i32>, mut potions: Vec<i32>, success: i64) -> Vec<i32> {
        potions.sort_unstable();
        let last = potions[potions.len() - 1] as i64;
        for x in spells.iter_mut() {
            let target = (success - 1) / *x as i64;
            if target < last { // 防止 i64 转成 i32 截断（这样不需要把 potions 中的数转成 i64 比较）
                let j = potions.partition_point(|&x| x <= target as i32);
                *x = (potions.len() - j) as i32;
            } else {
                *x = 0;
            }
        }
        spells
    }
}

复杂度分析

时间复杂度：$\mathcal{O}((n+m)\log m)$，其中 $n$ 为 $\textit{spells}$ 的长度，$m$ 为 $\textit{potions}$ 的长度。排序 $\mathcal{O}(m\log m)$。二分 $n$ 次，每次 $\mathcal{O}(\log m)$。
空间复杂度：$\mathcal{O}(1)$。忽略排序的栈开销，仅用到若干额外变量。

方法二：计数 + 值域后缀和

方法一得出的结论是，统计满足 $\textit{potions}[j]\ge \left\lfloor\dfrac{\textit{success}-1}{\textit{spell}[i]}\right\rfloor + 1$ 的药水的个数。

比如 $\textit{potions}=[1,2,2,3,5,5,5]$，要计算 $\ge 2$ 的药水的个数，我们可以统计每个数出现了多少次，记在一个 $\textit{cnt}$ 数组中。在这个例子中，$\textit{cnt}=[0,1,2,1,0,3]$，比如 $\textit{cnt}[5]=3$ 表示 $5$ 出现了 $3$ 次。

那么计算 $\textit{cnt}[2] + \textit{cnt}[3] + \textit{cnt}[4] + \textit{cnt}[5]=2+1+0+3=6$，就是 $\ge 2$ 的药水的个数。

但这样太慢了。如何加速？

借鉴前缀和的思想，我们可以倒着遍历 $\textit{cnt}$，原地计算 $\textit{cnt}$ 的后缀和，把 $\textit{cnt}[i]$ 更新为 $\ge i$ 的药水的个数。上面的 $\textit{cnt}$ 可以更新为 $[7,7,6,4,3,3]$。比如 $\textit{cnt}[2] = 6$ 表示 $\ge 2$ 的药水的个数。

class Solution:
    def successfulPairs(self, spells: List[int], potions: List[int], success: int) -> List[int]:
        mx = max(potions)
        cnt = [0] * (mx + 1)
        for y in potions:
            cnt[y] += 1  # 统计每种药水的出现次数

        # 计算 cnt 的后缀和
        for i in range(mx - 1, -1, -1):
            cnt[i] += cnt[i + 1]
        # 计算完毕后，cnt[i] 就是 potions 值 >= i 的药水个数

        for i, x in enumerate(spells):
            low = (success - 1) // x + 1
            spells[i] = cnt[low] if low <= mx else 0
        return spells

class Solution {
    public int[] successfulPairs(int[] spells, int[] potions, long success) {
        int mx = 0;
        for (int y : potions) {
            mx = Math.max(mx, y);
        }

        int[] cnt = new int[mx + 1];
        for (int y : potions) {
            cnt[y]++; // 统计每种药水的出现次数
        }

        // 计算 cnt 的后缀和
        for (int i = mx - 1; i >= 0; i--) {
            cnt[i] += cnt[i + 1];
        }
        // 计算完毕后，cnt[i] 就是 potions 值 >= i 的药水个数

        for (int i = 0; i < spells.length; i++) {
            long low = (success - 1) / spells[i] + 1;
            spells[i] = low <= mx ? cnt[(int) low] : 0;
        }
        return spells;
    }
}

class Solution {
public:
    vector<int> successfulPairs(vector<int>& spells, vector<int>& potions, long long success) {
        int mx = ranges::max(potions);
        vector<int> cnt(mx + 1);
        for (int y : potions) {
            cnt[y]++; // 统计每种药水的出现次数
        }

        // 计算 cnt 的后缀和
        for (int i = mx - 1; i >= 0; i--) {
            cnt[i] += cnt[i + 1];
        }
        // 计算完毕后，cnt[i] 就是 potions 值 >= i 的药水个数

        for (int& x : spells) {
            long long low = (success - 1) / x + 1;
            x = low <= mx ? cnt[low] : 0;
        }
        return spells;
    }
};

#define MAX(a, b) ((b) > (a) ? (b) : (a))

int* successfulPairs(int* spells, int spellsSize, int* potions, int potionsSize, long long success, int* returnSize) {
    int mx = 0;
    for (int i = 0; i < potionsSize; i++) {
        mx = MAX(mx, potions[i]);
    }

    int* cnt = calloc(mx + 1, sizeof(int));
    for (int i = 0; i < potionsSize; i++) {
        cnt[potions[i]]++; // 统计每种药水的出现次数
    }

    // 计算 cnt 的后缀和
    for (int i = mx - 1; i >= 0; i--) {
        cnt[i] += cnt[i + 1];
    }
    // 计算完毕后，cnt[i] 就是 potions 值 >= i 的药水个数

    for (int i = 0; i < spellsSize; i++) {
        long long low = (success - 1) / spells[i] + 1;
        spells[i] = low <= mx ? cnt[low] : 0;
    }

    *returnSize = spellsSize;
    free(cnt);
    return spells;
}

func successfulPairs(spells, potions []int, success int64) []int {
mx := slices.Max(potions)
cnt := make([]int, mx+1)
for _, y := range potions {
cnt[y]++ // 统计每种药水的出现次数
}

// 计算 cnt 的后缀和
for i := mx - 1; i >= 0; i-- {
cnt[i] += cnt[i+1]
}
// 计算完毕后，cnt[i] 就是 potions 值 >= i 的药水个数

for i, x := range spells {
low := (int(success)-1)/x + 1
if low <= mx {
spells[i] = cnt[low]
} else {
spells[i] = 0
}
}
return spells
}

var successfulPairs = function(spells, potions, success) {
    const mx = Math.max(...potions);
    const cnt = Array(mx + 1).fill(0);
    for (const y of potions) {
        cnt[y]++; // 统计每种药水的出现次数
    }

    // 计算 cnt 的后缀和
    for (let i = mx - 1; i >= 0; i--) {
        cnt[i] += cnt[i + 1];
    }
    // 计算完毕后，cnt[i] 就是 potions 值 >= i 的药水个数

    for (let i = 0; i < spells.length; i++) {
        const low = Math.ceil(success / spells[i]);
        spells[i] = low <= mx ? cnt[low] : 0;
    }
    return spells;
};

impl Solution {
    pub fn successful_pairs(mut spells: Vec<i32>, potions: Vec<i32>, success: i64) -> Vec<i32> {
        let mx = *potions.iter().max().unwrap() as usize;
        let mut cnt = vec![0; mx + 1];
        for y in potions {
            cnt[y as usize] += 1; // 统计每种药水的出现次数
        }

        // 计算 cnt 的后缀和
        for i in (0..mx).rev() {
            cnt[i] += cnt[i + 1];
        }
        // 计算完毕后，cnt[i] 就是 potions 值 >= i 的药水个数

        let success = success as usize;
        for x in spells.iter_mut() {
            let low = (success - 1) / *x as usize + 1;
            *x = if low <= mx { cnt[low] } else { 0 };
        }
        spells
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(n+m+U)$，其中 $n$ 为 $\textit{spells}$ 的长度，$m$ 为 $\textit{potions}$ 的长度，$U=\max(\textit{potions})$。
空间复杂度：$\mathcal{O}(U)$。

思考题

把乘法改成异或要怎么做？

这题是 1803. 统计异或值在范围内的数对有多少，做法见我的题解。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

LeetCode 每日一题题解
贪心 + 有序集合 / 并查集（Python/Java/C++/Go/JS/Rust）endlesscheng
2025年10月2日 13:27

贪心 + 有序集合 / 并查集（Python/Java/C++/Go/JS/Rust）

LeetCode 每日一题题解

作者 endlesscheng

2025年10月2日 13:27

第一个错误的思路

从左到右遍历 $\textit{rains}$：

如果 $\textit{rains}[i]=0$，把 $i$ 加到一个列表中，后面要抽水的时候再使用。
如果 $\textit{rains}[i]>0$ 且 $\textit{rains}[i]$ 是满的，随便取出列表中的一天用来抽水。如果列表是空的，那么必定会发生洪水。

错误原因：比如 $\textit{rains}=[1,0,0,1,1]$，抽水之后连着下了两天的雨，最后一天必定发洪水。注意抽水必须在两个 $1$ 之间，第二个 $0$ 无法用来抽干湖水。但按照上面的算法，我们会错误地认为第二个 $0$ 可以把湖 $1$ 抽干。

第二个错误的思路

从左到右遍历 $\textit{rains}$：

如果 $\textit{rains}[i]=0$，把 $i$ 入栈。
如果 $\textit{rains}[i]>0$ 且 $\textit{rains}[i]$ 是满的，那么用栈顶的那一天抽水。如果栈顶那一天比 $\textit{rains}[i]$ 装满的那一天还要早，那么必定会发生洪水。

错误原因：例如 $\textit{rains}=[1,0,2,0,1,2]$，如果第二个 $0$ 用来抽湖 $1$，那么湖 $2$ 就没法抽干。

正确思路

从左到右遍历 $\textit{rains}$：

如果 $\textit{rains}[i]=0$，把 $i$ 加到一个有序集合中。
如果 $\textit{rains}[i]>0$ 且 $\textit{rains}[i]$ 是满的。设 $\textit{rains}[i]$ 装满的那一天为 $j$，从有序集合中选出大于 $j$ 的最早的抽水日。

解释：越晚的抽水日，灵活性越大，可以用于更晚装满的湖。所以越晚的抽水日越应该留到后面再使用。

写法一：有序集合

###py

class Solution:
    def avoidFlood(self, rains: List[int]) -> List[int]:
        n = len(rains)
        ans = [-1] * n
        full_day = {}  # lake -> 装满日
        dry_day = SortedList()  # 未被使用的抽水日
        for i, lake in enumerate(rains):
            if lake == 0:
                ans[i] = 1  # 先随便选一个湖抽干
                dry_day.add(i)  # 保存抽水日
                continue
            if lake in full_day:
                j = full_day[lake]
                # 必须在 j 之后，i 之前把 lake 抽干
                # 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
                k = dry_day.bisect_right(j)
                if k == len(dry_day):
                    return []  # 无法阻止洪水
                d = dry_day[k]
                ans[d] = lake
                dry_day.discard(d)  # 移除已使用的抽水日
            full_day[lake] = i  # 插入或更新装满日
        return ans

###java

class Solution {
    public int[] avoidFlood(int[] rains) {
        int n = rains.length;
        int[] ans = new int[n];
        Map<Integer, Integer> fullDay = new HashMap<>(); // lake -> 装满日
        TreeSet<Integer> dryDay = new TreeSet<>(); // 未被使用的抽水日
        for (int i = 0; i < n; i++) {
            int lake = rains[i];
            if (lake == 0) {
                ans[i] = 1; // 先随便选一个湖抽干
                dryDay.add(i); // 保存抽水日
                continue;
            }
            Integer j = fullDay.get(lake);
            if (j != null) {
                // 必须在 j 之后，i 之前把 lake 抽干
                // 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
                Integer d = dryDay.higher(j);
                if (d == null) {
                    return new int[]{}; // 无法阻止洪水
                }
                ans[d] = lake;
                dryDay.remove(d); // 移除已使用的抽水日
            }
            ans[i] = -1;
            fullDay.put(lake, i); // 插入或更新装满日
        }
        return ans;
    }
}

###cpp

class Solution {
public:
    vector<int> avoidFlood(vector<int>& rains) {
        int n = rains.size();
        vector<int> ans(n, -1);
        unordered_map<int, int> full_day; // lake -> 装满日
        set<int> dry_day; // 未被使用的抽水日
        for (int i = 0; i < n; i++) {
            int lake = rains[i];
            if (lake == 0) {
                ans[i] = 1; // 先随便选一个湖抽干
                dry_day.insert(i); // 保存抽水日
                continue;
            }
            if (full_day.contains(lake)) {
                int j = full_day[lake];
                // 必须在 j 之后，i 之前把 lake 抽干
                // 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
                auto it = dry_day.upper_bound(j);
                if (it == dry_day.end()) {
                    return {}; // 无法阻止洪水
                }
                ans[*it] = lake;
                dry_day.erase(it); // 移除已使用的抽水日
            }
            full_day[lake] = i; // 插入或更新装满日
        }
        return ans;
    }
};

###go

// import "github.com/emirpasic/gods/v2/trees/redblacktree"
func avoidFlood(rains []int) []int {
ans := make([]int, len(rains))
fullDay := map[int]int{} // lake -> 装满日
dryDay := redblacktree.New[int, struct{}]() // 未被使用的抽水日
for i, lake := range rains {
if lake == 0 {
ans[i] = 1 // 先随便选一个湖抽干
dryDay.Put(i, struct{}{}) // 保存抽水日
continue
}
if j, ok := fullDay[lake]; ok {
// 必须在 j 之后，i 之前把 lake 抽干
// 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
node, _ := dryDay.Ceiling(j)
if node == nil {
return nil // 无法阻止洪水
}
ans[node.Key] = lake
dryDay.Remove(node.Key) // 移除已使用的抽水日
}
ans[i] = -1
fullDay[lake] = i // 插入或更新装满日
}
return ans
}

###js

const { AvlTree } = require('datastructures-js');

var avoidFlood = function(rains) {
    const n = rains.length;
    const ans = Array(n).fill(-1);
    const fullDay = new Map(); // lake -> 装满日
    const dryDay = new AvlTree((a, b) => a - b); // 未被使用的抽水日
    for (let i = 0; i < n; i++) {
        const lake = rains[i];
        if (lake === 0) {
            ans[i] = 1; // 先随便选一个湖抽干
            dryDay.insert(i); // 保存抽水日
            continue;
        }
        const j = fullDay.get(lake);
        if (j !== undefined) {
            // 必须在 j 之后，i 之前把 lake 抽干
            // 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
            const node = dryDay.ceil(j);
            if (node === null) {
                return []; // 无法阻止洪水
            }
            ans[node.getValue()] = lake;
            dryDay.removeNode(node); // 移除已使用的抽水日
        }
        fullDay.set(lake, i); // 插入或更新装满日
    }
    return ans;
};

###rust

use std::collections::{BTreeSet, HashMap};

impl Solution {
    pub fn avoid_flood(rains: Vec<i32>) -> Vec<i32> {
        let n = rains.len();
        let mut ans = vec![-1; n];
        let mut full_day = HashMap::new(); // lake -> 装满日
        let mut dry_day = BTreeSet::new(); // 未被使用的抽水日
        for (i, lake) in rains.into_iter().enumerate() {
            if lake == 0 {
                ans[i] = 1; // 先随便选一个湖抽干
                dry_day.insert(i); // 保存抽水日
                continue;
            }
            if let Some(&j) = full_day.get(&lake) {
                // 必须在 j 之后，i 之前把 lake 抽干
                // 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
                if let Some(&d) = dry_day.range(j..).next() {
                    ans[d] = lake;
                    dry_day.remove(&d); // 移除已使用的抽水日
                } else {
                    return vec![]; // 无法阻止洪水
                }
            }
            full_day.insert(lake, i); // 插入或更新装满日
        }
        ans
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(n\log n)$，其中 $n$ 是 $\textit{rains}$ 的长度。
空间复杂度：$\mathcal{O}(n)$。

写法二：并查集

删除 $i$ 时，用并查集把 $i$ 指向 $i+1$（或者 $\text{find}(i+1)$）。

例如删除 $2$，那么并查集中 $2\to 3$。

然后删除 $3$，那么并查集中 $2\to 3\to 4$。

查找 $\ge 2$ 的最小的未被删除的天，顺着并查集中的 $2$，可以找到 $4$，即为 $\ge 2$ 的最小的未被删除的天。

对于 $\textit{rains}[i]>0$ 的天，直接把 $i$ 删除。

注：关于并查集的完整模板，见数据结构题单。

###py

class Solution:
    def avoidFlood(self, rains: List[int]) -> List[int]:
        n = len(rains)
        # 非递归并查集
        fa = list(range(n + 1))
        def find(x: int) -> int:
            rt = x
            while fa[rt] != rt:
                rt = fa[rt]
            while fa[x] != rt:
                fa[x], x = rt, fa[x]
            return rt

        ans = [-1] * n
        full_day = {}  # lake -> 装满日
        for i, lake in enumerate(rains):
            if lake == 0:
                ans[i] = 1  # 先随便选一个湖抽干
                continue
            if lake in full_day:
                j = full_day[lake]
                # 必须在 j 之后，i 之前把 lake 抽干
                # 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
                dry_day = find(j + 1)
                if dry_day >= i:
                    return []  # 无法阻止洪水
                ans[dry_day] = lake
                fa[dry_day] = find(dry_day + 1)  # 删除 dry_day
            fa[i] = i + 1  # 删除 i
            full_day[lake] = i  # 插入或更新装满日
        return ans

###java

class Solution {
    public int[] avoidFlood(int[] rains) {
        int n = rains.length;
        int[] fa = new int[n + 1];
        for (int i = 0; i <= n; i++) {
            fa[i] = i;
        }

        int[] ans = new int[n];
        Map<Integer, Integer> fullDay = new HashMap<>(); // lake -> 装满日
        for (int i = 0; i < n; i++) {
            int lake = rains[i];
            if (lake == 0) {
                ans[i] = 1; // 先随便选一个湖抽干
                continue;
            }
            Integer j = fullDay.get(lake);
            if (j != null) {
                // 必须在 j 之后，i 之前把 lake 抽干
                // 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
                int dryDay = find(j + 1, fa);
                if (dryDay >= i) {
                    return new int[]{}; // 无法阻止洪水
                }
                ans[dryDay] = lake;
                fa[dryDay] = find(dryDay + 1, fa); // 删除 dryDay
            }
            ans[i] = -1;
            fa[i] = i + 1; // 删除 i
            fullDay.put(lake, i); // 插入或更新装满日
        }
        return ans;
    }

    private int find(int x, int[] fa) {
        if (fa[x] != x) {
            fa[x] = find(fa[x], fa);
        }
        return fa[x];
    }
}

###cpp

class Solution {
    vector<int> fa;

    int find(int x) {
        if (fa[x] != x) {
            fa[x] = find(fa[x]);
        }
        return fa[x];
    }

public:
    vector<int> avoidFlood(vector<int>& rains) {
        int n = rains.size();
        fa.resize(n + 1);
        ranges::iota(fa, 0); // 并查集初始化

        vector<int> ans(n, -1);
        unordered_map<int, int> full_day; // lake -> 装满日
        for (int i = 0; i < n; i++) {
            int lake = rains[i];
            if (lake == 0) {
                ans[i] = 1; // 先随便选一个湖抽干
                continue;
            }
            if (full_day.count(lake)) {
                int j = full_day[lake];
                // 必须在 j 之后，i 之前把 lake 抽干
                // 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
                int dry_day = find(j + 1);
                if (dry_day >= i) {
                    return {}; // 无法阻止洪水
                }
                ans[dry_day] = lake;
                fa[dry_day] = find(dry_day + 1); // 删除 dry_day
            }
            fa[i] = i + 1; // 删除 i
            full_day[lake] = i; // 插入或更新装满日
        }
        return ans;
    }
};

###go

func avoidFlood(rains []int) []int {
n := len(rains)
// 非递归并查集
fa := make([]int, n+1)
for i := range fa {
fa[i] = i
}
find := func(x int) int {
rt := x
for fa[rt] != rt {
rt = fa[rt]
}
for fa[x] != rt {
fa[x], x = rt, fa[x]
}
return rt
}

ans := make([]int, n)
fullDay := map[int]int{} // lake -> 装满日
for i, lake := range rains {
if lake == 0 {
ans[i] = 1 // 先随便选一个湖抽干
continue
}
if j, ok := fullDay[lake]; ok {
// 必须在 j 之后，i 之前把 lake 抽干
// 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
dryDay := find(j + 1)
if dryDay >= i {
return nil // 无法阻止洪水
}
ans[dryDay] = lake
fa[dryDay] = find(dryDay + 1) // 删除 dryDay
}
ans[i] = -1
fa[i] = i + 1 // 删除 i
fullDay[lake] = i // 插入或更新装满日
}
return ans
}

###js

var avoidFlood = function(rains) {
    const n = rains.length;
    const fa = Array(n + 1).fill(0).map((_, i) => i);

    function find(x) {
        if (fa[x] !== x) {
            fa[x] = find(fa[x]);
        }
        return fa[x];
    }

    const ans = Array(n).fill(-1);
    const fullDay = new Map(); // lake -> 装满日
    for (let i = 0; i < n; i++) {
        const lake = rains[i];
        if (lake === 0) {
            ans[i] = 1; // 先随便选一个湖抽干
            continue;
        }
        const j = fullDay.get(lake);
        if (j !== undefined) {
            // 必须在 j 之后，i 之前把 lake 抽干
            // 选一个最早的未被使用的抽水日，如果选晚的，可能会导致其他湖没有可用的抽水日
            const dryDay = find(j + 1);
            if (dryDay >= i) {
                return []; // 无法阻止洪水
            }
            ans[dryDay] = lake;
            fa[dryDay] = find(dryDay + 1); // 删除 dryDay
        }
        fa[i] = i + 1; // 删除 i
        fullDay.set(lake, i); // 插入或更新装满日
    }
    return ans;
};

###rust

use std::collections::HashMap;

impl Solution {
    pub fn avoid_flood(rains: Vec<i32>) -> Vec<i32> {
        let n = rains.len();
        let mut fa = (0..=n).collect::<Vec<_>>();

        fn find(x: usize, fa: &mut [usize]) -> usize {
            if fa[x] != x {
                fa[x] = find(fa[x], fa);
            }
            fa[x]
        }

        let mut ans = vec![-1; n];
        let mut full_day = HashMap::new(); // lake -> 装满日
        for (i, lake) in rains.into_iter().enumerate() {
            if lake == 0 {
                ans[i] = 1; // 先随便选一个湖抽干
                continue;
            }
            if let Some(&j) = full_day.get(&lake) {
                // 必须在 j 之后，i 之前把 lake 抽干
                let dry_day = find(j + 1, &mut fa);
                if dry_day >= i {
                    return vec![]; // 无法阻止洪水
                }
                ans[dry_day] = lake;
                fa[dry_day] = find(dry_day + 1, &mut fa); // 删除 dry_day
            }
            fa[i] = i + 1; // 删除 i
            full_day.insert(lake, i); // 插入或更新装满日
        }
        ans
    }
}

复杂度分析

时间复杂度：$\mathcal{O}(n)$，其中 $n$ 是 $\textit{rains}$ 的长度。本题并查集的操作是均摊 $\mathcal{O}(1)$。
空间复杂度：$\mathcal{O}(n)$。

专题训练

见下面数据结构题单的「§7.4 数组上的并查集」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

普通视图

方法一：懒删除堆

复杂度分析

方法二：倒序处理 + 维护最小值

复杂度分析

相似题目

专题训练

分类题单

复杂度分析

专题训练

分类题单

复杂度分析

专题训练

分类题单

复杂度分析

专题训练

分类题单

位运算

复杂度分析

数学

复杂度分析

分类题单

复杂度分析

分类题单

复杂度分析

分类题单

方法一：枚举

复杂度分析

方法二：倒序贪心

复杂度分析

相似题目

专题训练

分类题单

方法一：差分

答疑

复杂度分析

方法二：同向三指针 + 同向双指针

核心思路

同向三指针

同向双指针

小优化

复杂度分析

专题训练

分类题单

不轮转

枚举轮转到最左边的下标

复杂度分析

分类题单

优化

复杂度分析

专题训练

分类题单

写法一

复杂度分析

写法二

复杂度分析

相似题目

分类题单

复杂度分析

专题训练

分类题单

k = 1

k = 2

一般情况

写法一

写法二

复杂度分析

专题训练

分类题单

方法一：正反两次扫描

复杂度分析

方法二：递推开始时间

复杂度分析

方法三：record 优化

复杂度分析（最坏情况）

复杂度分析（平均情况）

方法四：凸包 + 二分

复杂度分析

分类题单

方法一：排序 + 二分查找