【单调栈】2024D-找朋友
【单调栈】2024D-找朋友
[P2650] 【单调栈】2024D-找朋友
本题练习地址：https://www.algomooc.com/problem/P2650
题目描述
在学校中，N个小朋友站成一队， 第i个小朋友的身高为height[i]，第i个小朋友可以看到的右边的第一个比自己身高更高的小朋友j，那么j是i的好朋友（j > i）。请重新生成一个列表，对应位置的输出是每个小朋友的好朋友位置，如果没有看到好朋友，请在该位置用0代替。小朋友人数范围是 [0, 40000]。
输入描述
第一行输入N，表示有N个小朋友
第二行输入N个小朋友的身高height[i]，都是整数
输出描述
输出N个小朋友的好朋友的位置
示例一
输入
2
100 95输出
0 0示例二
输入
8
123 124 125 121 119 122 126 123输出
1 2 6 5 5 6 0 0解题思路
两种写法的比较
注意，本题和LC739. 每日温度非常类似。区别在于，本题需要找到的是右边下一个更大元素的索引，而非与当前元素的间隔，显然变得更加简单了。
我们讲过，类似这种要求寻找左边/右边最近的更大/更小元素的题目，均可以使用单调栈来完成。
对于单调栈的题目，既可以正序遍历也可以逆序遍历数组来完成，重点在于理解单调栈的原理，同学们只需要选择适合自己理解的方法来完成即可。以下表格总结了两种不同遍历顺序的异同点。
| 正序遍历 | 逆序遍历 | |
|---|---|---|
| 单调栈顺序 | 栈中储存的索引所对应在原数组中的元素大小，从栈底至栈顶单调递减，即更大的数（的下标）位于栈底 | |
| 入栈时机 | 栈顶元素反复出栈并修改ans之后，进行入栈。且栈中元素为下标i，而非身高h | |
修改ans时机 | i为preIndex的下一个更大元素的下标，在出栈过程中，即在while内修改ans[preIndex] | stack[-1]为i的下一个更大元素的下标，在出栈结束后，即在while外修改ans[i] |
| 出栈条件 | 包含h > height[stack[-1]] | 包含h >= height[stack[-1]] |
正序遍历图解
正序遍历原数组，初始化单调栈为空，答案数组均为0。
暂时无法在飞书文档外展示此内容
遍历到nums[0]，单调栈中没有元素，索引0入栈。
暂时无法在飞书文档外展示此内容
继续遍历到到nums[1]，栈顶索引0对应的123小于当前元素124。
栈顶元素0出栈，同时修改ans[0]为当前索引1。
出栈和修改结束后，当前索引1入栈。
对于正序遍历写法而言：修改答案是在出栈的时候进行的。
暂时无法在飞书文档外展示此内容
遍历到nums[2]，过程也是类似的。
暂时无法在飞书文档外展示此内容
继续遍历nums[3]和nums[4]，由于121和119相较于125都是递减的，因此直接入栈，无需修改ans。
暂时无法在飞书文档外展示此内容
继续遍历nums[5]，当前元素大于栈顶索引4和3对应的元素的119和121。
将它们弹出，同时依次修改ans[4]和ans[3]为当前索引5。
出栈和修改结束后，当前索引5入栈。
暂时无法在飞书文档外展示此内容
继续遍历nums[6]，当前元素大于栈顶索引5和2对应的元素的122和125。
将它们弹出，同时依次修改ans[5]和ans[2]为当前索引6。
出栈和修改结束后，当前索引6入栈。
暂时无法在飞书文档外展示此内容
最后遍历到nums[7]，小于栈顶索引对应的元素126，无需弹出也无需修改答案数组。直接入栈。
暂时无法在飞书文档外展示此内容
最终答案数组的结果即为答案。
因此正序遍历原数组的单调栈算法的整体框架如下
# 正序遍历写法
for i, h in enumerate(height):
# 先while判断出栈
# 在while中修改ans
while len(stack) > 0 and h > height[stack[-1]]:
preIndex = stack.pop()
ans[preIndex] = i
# 后入栈
stack.append(i)逆序遍历图解
逆序遍历原数组，初始化单调栈为空，答案数组均为0。
暂时无法在飞书文档外展示此内容
遍历到到nums[7]，单调栈中没有元素，索引7入栈。
暂时无法在飞书文档外展示此内容
继续遍历到nums[6]，栈顶索引7对应的123小于等于当前元素126。
即使遍历到6更前面的位置（比如索引5），126必然会挡住后面的这个123，故123此时已经无用了。
逆序遍历写法的出栈操作其实隐含了贪心思想
栈中的索引7出栈。出栈后栈中没有元素，说明不存在位于右边的比126更大的值，无需修改答案数组，直接将当前索引6入栈。
暂时无法在飞书文档外展示此内容
继续遍历到nums[5]。由于122小于栈顶元素对应的126，可知此时栈顶索引6就是当前元素122右边第一个最大的元素的索引，修改ans[5]。
暂时无法在飞书文档外展示此内容
同理，继续遍历到nums[4]。由于119小于栈顶元素对应的122，可知此时栈顶索引5就是当前元素119右边第一个最大的元素的索引，修改ans[4]。而后将索引4入栈。
暂时无法在飞书文档外展示此内容
继续遍历到nums[3]。通过while循环，将栈顶所有对应元素小于等于当前元素121的索引弹出（即119）。
全部弹出后，此时栈顶仍然存在索引5，即为当前元素121右边的第一个更大值对应的索引。
修改ans[3]。而后将索引3入栈。
暂时无法在飞书文档外展示此内容
继续遍历到nums[2]。通过while循环，将栈顶所有对应元素小于等于当前元素125的索引弹出（即121、122）。
全部弹出后，此时栈顶仍然存在索引6，即为当前元素125右边的第一个更大值对应的索引。
修改ans[3]。而后将索引2入栈。
暂时无法在飞书文档外展示此内容
继续依次遍历nums[1]和nums[0]，其实过程和遍历到nums[5]和nums[4]是类似的，修改ans[1]和ans[0]，最终得到答案数组。
因此逆序遍历原数组的单调栈算法的整体框架如下
# 逆序遍历写法
for i in range(n-1, -1, -1):
h = height[i]
# 先while判断出栈
while len(stack) > 0 and h >= height[stack[-1]]:
stack.pop()
# 在while外修改ans
if len(stack) > 0:
# ans[i]修改为stack[-1]
ans[i] = stack[-1]
# 后入栈
stack.append(i)思考
作为一道经典的单调栈题目，我们现在已经掌握了其正序和逆序两种写法。
这道题是非常值得反复琢磨的题目。
请课后或者完成本题后再次思考以下问题：
- 为什么正序写法
while的条件是h > height[stack[-1]]，而逆序写法的是h >= height[stack[-1]] - 如果改成找右边第一个更小元素，代码应该如何修改？
- 如果改成找左边第一个更大元素，代码应该如何修改？
- 如果改成找左边第一个更小元素，代码应该如何修改？
代码
解法一
Python
正序遍历height构建单调栈。
# 题目：2024D-找朋友
# 分值：100
# 作者：许老师-闭着眼睛学数理化
# 算法：单调栈-正序遍历原数组
# 代码看不懂的地方，请直接在群上提问
# 输入小朋友个数n
n = int(input())
# 输入N个小朋友的高度数组
height = list(map(int, input().split()))
# 构建一个单调栈，用来存放不同小朋友的身高的索引
# 栈中储存的索引所对应在height中的元素大小，从栈底至栈顶单调递减
# 即更大的数（的下标）位于栈底
stack = list()
# 构建列表ans，用来保存输出结果
# 初始化其中所有的元素均为0
ans = [0] * n
# 从头开始遍历每一个小朋友的身高
for i, h in enumerate(height):
# 第i个小朋友的身高h，需要不断地与栈顶元素比较
# 如果栈顶元素存在并且h【大于】栈顶元素stack[-1]
# 意味着栈顶元素找到了右边最近的比他更高的身高h
while len(stack) > 0 and h > height[stack[-1]]:
# 首先获取栈顶元素的值，也就是上一个比h小的身高的索引值
preIndex = stack.pop()
# i即为preIndex这个索引所对应的，下一个最近身高
ans[preIndex] = i
# 再把当前小朋友身高的下标i存放到栈中
# 注意：所储存的是下标i，而不是身高h
stack.append(i)
# ans中的int元素转成str后才能合并成字符串
print(" ".join(map(str, ans)))Java
import java.util.Scanner;
import java.util.Stack;
public class Main {
public static void main(String[] args) {
Scanner scanner = new Scanner(System.in);
int n = scanner.nextInt();
int[] height = new int[n];
for (int i = 0; i < n; i++) {
height[i] = scanner.nextInt();
}
Stack<Integer> stack = new Stack<>();
int[] ans = new int[n];
// 从头开始遍历每一个小朋友的身高
for (int i = 0; i < n; i++) {
int h = height[i];
// 第i个小朋友的身高h，需要不断地与栈顶元素比较
// 如果栈顶元素存在并且h > 栈顶元素 stack.peek()
// 意味着栈顶元素找到了右边最近的比他更高的身高h
while (!stack.isEmpty() && h > height[stack.peek()]) {
// 首先获取栈顶元素的值，也就是上一个比h小的身高的索引值
int preIndex = stack.pop();
// i即为preIndex这个索引所对应的，下一个最近身高
ans[preIndex] = i;
}
// 再把当前小朋友身高的下标i存放到栈中
stack.push(i);
}
// ans中的int元素转成str后才能合并成字符串
for (int i = 0; i < n; i++) {
System.out.print(ans[i] + " ");
}
System.out.println();
}
}C++
#include <iostream>
#include <sstream>
#include <vector>
#include <stack>
using namespace std;
int main() {
int n;
cin >> n;
cin.ignore();
vector<int> height(n);
for (int i = 0; i < n; i++) {
cin >> height[i];
}
stack<int> stk;
vector<int> ans(n, 0);
for (int i = 0; i < n; i++) {
int h = height[i];
while (!stk.empty() && h > height[stk.top()]) {
int preIndex = stk.top();
stk.pop();
ans[preIndex] = i;
}
stk.push(i);
}
for (int i = 0; i < n; i++) {
cout << ans[i] << " ";
}
cout << endl;
return 0;
}解法二
逆序遍历height构建单调栈。
Python
# 题目：2024D-找朋友
# 分值：100
# 作者：许老师-闭着眼睛学数理化
# 算法：单调栈-逆序遍历原数组
# 代码看不懂的地方，请直接在群上提问
# 输入小朋友个数n
n = int(input())
# 输入N个小朋友的高度数组
height = list(map(int, input().split()))
# 构建一个单调栈，用来存放不同小朋友的身高的索引
# 栈中储存的索引所对应在height中的元素大小，从栈底至栈顶单调递增
# 即更大的数（的下标）位于栈底
stack = list()
# 构建列表ans，用来保存输出结果
# 初始化其中所有的元素均为0
ans = [0] * n
# 逆序遍历每一个小朋友的身高
for i in range(n-1, -1, -1):
h = height[i]
# 第i个小朋友的身高h，需要不断地与栈顶元素比较
# 如果栈顶元素存在并且h【大于等于】栈顶元素stack[-1]
# 说明栈顶元素stack[-1]并不是身高h右边最近的比h更大的元素
# 需要将栈顶元素弹出，继续寻找比h大的栈顶元素
while len(stack) > 0 and h >= height[stack[-1]]:
# 栈顶元素下标对应的身高不大于当前身高h，不是符合要求的更大身高，弹出
stack.pop()
# 完成弹出后，如果栈顶仍存在元素，说明stack[-1]所对应的身高，是严格比h大的下一个身高
if len(stack) > 0:
# ans[i]修改为stack[-1]
ans[i] = stack[-1]
# 再把当前小朋友身高的下标i存放到栈中
# 注意：所储存的是下标i，而不是身高h
stack.append(i)
# ans中的int元素转成str后才能合并成字符串
print(" ".join(map(str, ans)))Java
import java.util.Scanner;
import java.util.Stack;
public class Main {
public static void main(String[] args) {
Scanner scanner = new Scanner(System.in);
int n = scanner.nextInt();
int[] height = new int[n];
for (int i = 0; i < n; i++) {
height[i] = scanner.nextInt();
}
Stack<Integer> stack = new Stack<>();
int[] ans = new int[n];
// 逆序遍历每一个小朋友的身高
for (int i = n - 1; i >= 0; i--) {
int h = height[i];
// 第i个小朋友的身高h，需要不断地与栈顶元素比较
// 如果栈顶元素存在并且h >= 栈顶元素 stack.peek()
// 说明栈顶元素 stack.peek() 并不是身高h右边最近的比h更大的元素
// 需要将栈顶元素弹出，继续寻找比h大的栈顶元素
while (!stack.isEmpty() && h >= height[stack.peek()]) {
// 栈顶元素下标对应的身高不大于当前身高h，不是符合要求的更大身高，弹出
stack.pop();
}
// 完成弹出后，如果栈顶仍存在元素，说明 stack.peek() 所对应的身高，是严格比h大的下一个身高
if (!stack.isEmpty()) {
// ans[i] 修改为 stack.peek()
ans[i] = stack.peek();
}
// 再把当前小朋友身高的下标i存放到栈中
stack.push(i);
}
// ans中的int元素转成str后才能合并成字符串
for (int i = 0; i < n; i++) {
System.out.print(ans[i] + " ");
}
System.out.println();
}
}C++
#include <iostream>
#include <sstream>
#include <vector>
#include <stack>
using namespace std;
int main() {
int n;
cin >> n;
cin.ignore();
vector<int> height(n);
for (int i = 0; i < n; i++) {
cin >> height[i];
}
stack<int> stk;
vector<int> ans(n, 0);
for (int i = n - 1; i >= 0; i--) {
int h = height[i];
while (!stk.empty() && h >= height[stk.top()]) {
stk.pop();
}
if (!stk.empty()) {
ans[i] = stk.top();
}
stk.push(i);
}
for (int i = 0; i < n; i++) {
cout << ans[i] << " ";
}
cout << endl;
return 0;
}时空复杂度
时间复杂度：O(N)。不管是正序还是逆序遍历，均仅需一次遍历height数组。
空间复杂度：O(N)。单调栈所占用的额外空间。