[LeetCode in Python] 5381 (M) queries on a permutation with key 查询带键的排列

题目：

https://leetcode-cn.com/problems/queries-on-a-permutation-with-key/

给你一个待查数组 queries ，数组中的元素为 1 到 m 之间的正整数。
请你根据以下规则处理所有待查项 queries[i]（从 i=0 到 i=queries.length-1）：

一开始，排列 P=[1,2,3,...,m]。
对于当前的 i ，请你找出待查项 queries[i] 在排列 P 中的位置（下标从 0 开始），然后将其从原位置移动到排列 P 的起始位置（即下标为 0 处）。
注意， queries[i] 在 P 中的位置就是 queries[i] 的查询结果。
请你以数组形式返回待查数组? queries 的查询结果。

示例 1：

输入：queries = [3,1,2,1], m = 5
输出：[2,1,2,1]
解释：待查数组 queries 处理如下：
对于 i=0: queries[i]=3, P=[1,2,3,4,5], 3 在 P 中的位置是 2，接着我们把 3 移动到 P 的起始位置，得到 P=[3,1,2,4,5] 。
对于 i=1: queries[i]=1, P=[3,1,2,4,5], 1 在 P 中的位置是 1，接着我们把 1 移动到 P 的起始位置，得到 P=[1,3,2,4,5] 。
对于 i=2: queries[i]=2, P=[1,3,2,4,5], 2 在 P 中的位置是 2，接着我们把 2 移动到 P 的起始位置，得到 P=[2,1,3,4,5] 。
对于 i=3: queries[i]=1, P=[2,1,3,4,5], 1 在 P 中的位置是 1，接着我们把 1 移动到 P 的起始位置，得到 P=[1,2,3,4,5] 。
因此，返回的结果数组为 [2,1,2,1] 。

示例 2：

输入：queries = [4,1,2,2], m = 4
输出：[3,1,2,0]

示例 3：

输入：queries = [7,5,5,8,3], m = 8
输出：[6,5,0,7,5]

提示：

1 <= m <= 10^3
1 <= queries.length <= m
1 <= queries[i] <= m

仿真解题思路

• 最常见的思路是仿真。
• 太简单以至于没必要说什么。

代码

class Solution:
    def processQueries(self, queries: List[int], m: int) -> List[int]:
        a = [i for i in range(m+1)]
        d = {i:i for i in range(m+1)}
        
        res = []
        for q in queries:
            pos = d[q]
            res.append(pos-1)
            
            for i in range(pos,1,-1):
                a[i] = a[i-1]
                d[a[i]] = i
            
            a[1] = q
            d[q] = 1
        
        return res

Fenwick Tree 解题思路

本题答案来自awice，O(n logn)，我对其添加了一点点注释。

• Fenwick Tree可以用来简化前缀和的更新和计算。
• 不熟悉的可以参考花花酱的视频。
• 本题因为有限制1 <= queries.length <= m，所以思路是：
• 构造一个2*m长度的Fenwick Tree，
• 前面留出m的空白，
• 后面一半是真正的tree data，
• 使用一个字典来维护数字到位置的映射，
• 当需要移动某个元素到开头时，
• 这个开头的位置是从m开始，
• 然后是 m-1, m-2, ...
• 每次移动都会更新tree data及位置字典，
• tree data的更新，
• 包括元素之后全部元素的更新，
• 及首元素更新之后导致的全体元素更新。

代码

class Fenwick:
    def __init__(self, n):
        sz = 1
        while sz <= n:
            sz *= 2
        self.size = sz
        self.data = [0] * sz

    def sum(self, i):
        s = 0
        while i > 0:
            s += self.data[i]
            i -= i & -i
        return s

    def add(self, i, x):
        while i < self.size:
            self.data[i] += x
            i += i & -i

class Solution:
    def processQueries(self, queries: List[int], m: int) -> List[int]:
        # - because queries.length <= m, so use double space 
        fenw = Fenwick(2 * m)

        # - vimap keeps the position of value i
        vimap = {}
        for i in range(1, m + 1):
            fenw.add(i + m, 1)
            vimap[i] = i + m

        # - next head position
        cur = m
        
        ans = []
        for q in queries:
            # - get current position of value q
            i = vimap.pop(q)

            # - rank means index
            rank = fenw.sum(i-1)
            ans.append(rank)
            
            # - move q to the head position
            vimap[q] = cur

            # - all elements behind position i will move left, so rank--
            fenw.add(i, -1)

            # - all elements value shift +1
            fenw.add(cur, 1)

            # - move next head position to left
            cur -= 1
            
        return ans