华为OD机试 Python 【找朋友】

 题目

在学校中，N 个小朋友站成一队， 第 i 个小朋友的身高为height[i]，第 i 个小朋友可以看到的第一个比自己身高更高的小朋友 j ，那么 j 是 i 的好朋友(要求j > i)。

请输出一个数组，对应位置的输出是每个小朋友的好朋友位置，如果没有看到好朋友，则输出0。

输入描述

第一行输入N，N表示有N个小朋友

第二行输入N个小朋友的身高height[i]，都是整数

输出描述

输出N个小朋友的好朋友的位置

示例1：

输入    

2

100 95

输出    

0 0

示例2：

输入：

8
123 124 125 121 119 122 126 123

输出：

1 2 6 5 5 6 0 0

题意解读：

这个题目意思是，给定一列小朋友站成一队，每个小朋友有一个身高值。我们输入的就是这些小朋友的身高。对于队列中的每一个小朋友，我们要找出站在他后面的、并且比他更高的第一个小朋友的位置。

以 123 124 125 121 119 122 126 123 为例，他表示这八个小朋友的身高。

1. 第 0 个小朋友身高123：我们需要找到他的后面，第一个比他高的小朋友。于是我们找到了身高为 124 的小朋友。而 124 的索引为 2，所以 123 的好朋友是在位置 1。
2. 第 1 个小朋友身高124：向后看，第一个比他高的是125的小朋友，所以125的好朋友是在位置2。
3. 第 2 个小朋友身高125：向后看，直到找到126的小朋友，因为他是唯一一个比125高的，所以125的好朋友是在位置 6。
4. 后面的以此类推，对于每个小朋友，都往后找，找到第一个身高比他高的小朋友的索引，就能得出答案。

所以正确的答案是 [1, 2, 6, 5, 5, 6, 0, 0]。

解题思路

华为 OD 机试：找朋友【空间换时间】

暴力解法

通过两层循环遍历来找到每个小朋友的“好朋友”。第一层循环遍历每个小朋友，第二层循环从当前小朋友开始向后查找第一个比当前小朋友身高更高的小朋友。 因此时间复杂度是 n*n

def find_next_taller(heights):
    n = len(heights)
    result = [0] * n  
    
    # 对每个小朋友，向后查找第一个比他高的小朋友
    for i in range(n - 1): 
        for j in range(i + 1, n):
            if heights[j] > heights[i]:
                result[i] = j  
                break  
    
    return result

def main():
    # 小朋友的数量
    n = int(input())
    
    # 每个小朋友的身高
    heights = [int(x) for x in input().split()]
    
    # 找到下一个更高的小朋友
    result = find_next_taller(heights)
    
    for position in result:
        print(position, end=" ")

if __name__ == '__main__':
    main()

栈解法

开始时，栈是空的，然后依次遍历每个小朋友。

 ·   如果当前小朋友的身高超过栈顶小朋友的身高，则说明当前小朋友是栈顶小朋友的“好朋友”（即第一个比他高的小朋友）。在这种情况下，将栈顶元素弹出，并更新相应位置的输出数组，表示找到了该位置小朋友的好朋友。并将当前小朋友的位置入栈。

·  如果当前小朋友没有超过栈顶小朋友，则将当前小朋友的位置入栈。

·  重复以上步骤，直到遍历完所有小朋友。

遍历结束时，栈中剩余的元素表示那些没有找到比自己高的小朋友。根据题目要求，他们的好朋友就是 0

def get_friends_position(heights):
    stack = []  
    friends_position = [0] * len(heights)  
    
    for i in range(len(heights)):
        # 当栈不为空，且当前小朋友身高大于栈顶小朋友的身高时
        while stack and heights[i] > heights[stack[-1]]:
            # 更新栈顶小朋友的好朋友位置为当前小朋友的位置
            friends_position[stack.pop()] = i
        # 将当前小朋友的位置入栈
        stack.append(i)
    
    return friends_position

def main():
    # 小朋友的数量
    n = int(input())
    
    # 小朋友的身高
    heights = [int(x) for x in input().split()]
    
    result = get_friends_position(heights)
    
    for position in result:
        print(position, end=" ")

main()