Leetcode 3008. Find Beautiful Indices in the Given Array II

• Leetcode 3008. Find Beautiful Indices in the Given Array II
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：3008. Find Beautiful Indices in the Given Array II

1. 解题思路

这一题其实算是套路题了，知道的话就挺快的，不知道的话就会很难了……

思路上来说的话其实很直接，就是首先获取原始的字符串s当中所有的子串a和子串b出现的位置数组loc_a和loc_b，然后对比两个数组，找到所有loc_a当中的元素，使得loc_b当中存在某一元素与其距离的绝对值不大于k。

首先我们考察第二部分的问题，这个的话我们对于两个有序数组，使用滑动窗口就能在$O(N)$的时间复杂度内处理了，并不复杂，注意一下滑动窗口的移动条件就行了，这里就不过多展开了。

剩下的，就是第一部分的问题，如何在一个字符串s当中获取所有的子串sub出现的位置。

这个问题其实蛮难的，自己想的话估计能想死人，不过现在已经是个套路题了，用z算法就能够直接搞定了，我之前也写过一个博客（经典算法：Z算法（z algorithm））介绍过这个算法，里面有个例子一模一样，直接复制里面的代码来用就行了。如果觉得我写的不好的话也可以自己去wiki或者知乎搜搜，挺多的……

2. 代码实现

给出python代码实现如下：

def z_algorithm(s):
    n = len(s)
    z = [0 for _ in range(n)]
    l, r = -1, -1
    for i in range(1, n):
        if i > r:
            l, r = i, i
            while r < n and s[r-l] == s[r]:
                r += 1
            z[i] = r-l
            r -= 1
        else:
            k = i - l
            if z[k] < r - i + 1:
                z[i] = z[k]
            else:
                l = i
                while r < n and s[r-l] == s[r]:
                    r += 1
                z[i] = r-l
                r -= 1
    z[0] = n
    return z

class Solution:
    def beautifulIndices(self, s: str, a: str, b: str, k: int) -> List[int]:
        
        def find_all(s, sub):
            z = z_algorithm(sub+s)
            n, m = len(s), len(sub)
            z = z[m:]
            return [idx for idx, l in enumerate(z) if l >= m]
        
        alist = find_all(s, a)
        blist = find_all(s, b)

        ans = []
        i, j, n = 0, 0, len(blist)
        for idx in alist:
            while i < n and blist[i] < idx-k:
                i += 1
            while j < n and blist[j] <= idx+k:
                j += 1
            if j-i>0:
                ans.append(idx)
        return ans

提交代码评测得到：耗时1619ms，占用内存65MB。