【Leetcode】1562. 查找大小为 M 的最新分组——1928

题目

1562. 查找大小为 M 的最新分组

给你一个数组 arr ，该数组表示一个从 1 到 n 的数字排列。有一个长度为 n 的二进制字符串，该字符串上的所有位最初都设置为 0 。

在从 1 到 n 的每个步骤 i 中（假设二进制字符串和 arr 都是从 1 开始索引的情况下），二进制字符串上位于位置 arr[i] 的位将会设为 1 。
给你一个整数 m ，请你找出二进制字符串上存在长度为 m 的一组 1 的最后步骤。一组 1 是一个连续的、由 1 组成的子串，且左右两边不再有可以延伸的 1 。
返回存在长度 恰好 为 m 的 一组 1 的最后步骤。如果不存在这样的步骤，请返回 -1 。
提示：

• \(n == arr.length\)
• \(1 \leq n \leq 10^5\)
• \(1 \leq= arr[i]\leq= n\)
• \(arr\) 中的所有整数 互不相同
• \(1 \leq= m \leq= arr.length\)

思路

问题分析

问题是将n个位置按照要求进行变为1，求解的是在最后的那个步骤仍然存在最长连续长度为m的1。

由于要求的是一定是要去连续最长的一段，那么实际上对于某一个位置变为1之后，他可能发生的操作就是左边存在一段连续的1，右边也存在一段连续的1，之后这三段进行合并为一段长的。

如果将这个过程看作是集合之间的合并，那么久很容易想到了使用并查集，当前位置准备为1，那么如果左边有连续的1，就将其合并，右边存在就也将其合并。

并查集求解

在使用并查集的时候，需要额外注意的是，当每个位置没有被标记为1的时候，我们可以使用fa[i]=-1表示为其并不存在于并查集中。
此外为了方便起见fa的长度设置为n+2，便于对于每个位置-1和+1位置的操作判断。

n = len(arr)
fa = [-1 for i in range(n+2)] # 并查集
num = [1]*(n+1) # 记录连续1的长度
def find(x):
	if x != fa[x]: fa[x] = find(fa[x])
	return fa[x]

那么根据并查集我们就可以很容易的实现整个连续1的序列的集合的过程。
但是需要注意的点在于当前我们对于左边和右边合并完毕之后，如果其长度等于m，则表示当前位置可以得到m长度的连续1，但是当其不等于m的时候，我们就不能简单的判断其一定不存在m长度的连续1。

因此可以引入一个额外的队列，表示原有其可能存在m长度的连续1的位置。
当当前位置是m长度的连续1，我们将其放入队列
如果队列首部记录的位置中其通过并查集进行修改了头结点，代表其连续1长度一定发生变化，那么该节点弹出。
若队列中仍存在某位置为m长度的连续1，那么就可以在当前的操作中继续更新结果。
完整的代码如下：

class Solution:
    def findLatestStep(self, arr: List[int], m: int) -> int:
        n = len(arr)
        fa = [-1 for i in range(n+2)] # 并查集
        num = [1]*(n+1) # 记录连续1的长度
        ans = -1
        def find(x):
            if x != fa[x]: fa[x] = find(fa[x])
            return fa[x]
        q = []
        for i,x in enumerate(arr):
            fa[x] = x # x位置变成1，加入并查集
            for y in [x-1,x+1]:
                if fa[y]!=-1:
                    node = find(y)
                    num[x]+=num[node]
                    fa[node] = x
            if num[x] == m:
                q.append(x)
            while q and (find(q[-1])!=q[-1] or num[q[-1]]!=m): q.pop()
            if q: ans = max(ans, i+1)
        return ans

并查集优化

可以看到上述过程是使用了一个额外的空间存储一个队列以表示每个步骤之后是否有满足条件的，但是实际上相比于所有步骤的，我们其实更关心哪些相对步骤靠后的。

首先考虑第一种情况，当前步骤执行完，形成了一个包含当前位置的m长度的连续1，那么此时可以确保当前位置一定合理。

此外，我们对于整个数组的步骤执行情况出发，由于初始给定的数组是1~n的排列，那么无论如何到最后一定是每个位置都是1，那么对于恰好为m的 一组1的最后步骤的下一个步骤而言，他会将这个符合条件的给覆盖掉，因此，我们在覆盖之前可以进行判断，如果合并之前的左右两边存在满足条件的，那么实际上上一步就是恰好为m的一组1的步骤

根据上述分析，对于之前代码的q就可以直接去掉了。修改后的代码如下所示：

class Solution:
    def findLatestStep(self, arr: List[int], m: int) -> int:
        n = len(arr)
        fa = [-1 for i in range(n+2)] # 并查集
        num = [1]*(n+1) # 记录连续1的长度
        ans = -1
        def find(x):
            if x != fa[x]: fa[x] = find(fa[x])
            return fa[x]
        for i,x in enumerate(arr):
            fa[x] = x # x位置变成1，加入并查集
            for y in [x-1,x+1]:
                if fa[y]!=-1:
                    node = find(y)
                    if num[node] == m:
                        ans = i
                    num[x]+=num[node]
                    fa[node] = x
            if num[x] == m:
                ans = i+1
        return ans

进一步优化

之前的算法中，我们采用了并查集，这样就可以实现了每个位置都能进行查询其所在的分组或者叫集合。

然而实际上，在本题目中，我们可以看到的情况是，如果当前i~j已经全为1，那么实际上之后就不会再访问i+1~j-1之间的所有位置了，因此其余的所有位置与i+1~j-1都不相邻。那么就可以将上述的并查集记录的方式进行变换。只记录可能为之后某步骤的相邻可能组成新的集合的边界位置进行记录即可。

class Solution:
    def findLatestStep(self, arr: List[int], m: int) -> int:
        res = -1
        memo = {}
        for i, idx in enumerate(arr):
            l = memo.get(idx - 1, 0)
            r = memo.get(idx + 1, 0)
            cur = l + 1 + r
            if (l == m) or (r == m):
                res = i
            if (cur == m):
                res = i+1
            memo[idx - l] = cur
            memo[idx + r] = cur
        return res