leetcode 1562. 查找大小为 M 的最新分组 树状数组解法或并查集 巧妙

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

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

在从 1 到 n 的每个步骤 i 中（假设二进制字符串和 arr 都是从 1 开始索引的情况下），二进制字符串上位于位置 arr[i] 的位将会设为 1 。

给你一个整数 m ，请你找出二进制字符串上存在长度为 m 的一组 1 的最后步骤。一组 1 是一个连续的、由 1 组成的子串，且左右两边不再有可以延伸的 1 。

返回存在长度 恰好 为 m 的 一组 1 的最后步骤。如果不存在这样的步骤，请返回 -1 。

示例 1：

输入：arr = [3,5,1,2,4], m = 1
输出：4
解释：
步骤 1：“00100”，由 1 构成的组：[“1”]
步骤 2：“00101”，由 1 构成的组：[“1”, “1”]
步骤 3：“10101”，由 1 构成的组：[“1”, “1”, “1”]
步骤 4：“11101”，由 1 构成的组：[“111”, “1”]
步骤 5：“11111”，由 1 构成的组：[“11111”]
存在长度为 1 的一组 1 的最后步骤是步骤 4 。

示例 2：

输入：arr = [3,1,5,4,2], m = 2
输出：-1
解释：
步骤 1：“00100”，由 1 构成的组：[“1”]
步骤 2：“10100”，由 1 构成的组：[“1”, “1”]
步骤 3：“10101”，由 1 构成的组：[“1”, “1”, “1”]
步骤 4：“10111”，由 1 构成的组：[“1”, “111”]
步骤 5：“11111”，由 1 构成的组：[“11111”]
不管是哪一步骤都无法形成长度为 2 的一组 1 。

示例 3：

输入：arr = [1], m = 1
输出：1

示例 4：

输入：arr = [2,1], m = 2
输出：2

提示：

n == arr.length
1 <= n <= 10^5
1 <= arr[i] <= n
arr 中的所有整数 互不相同
1 <= m <= arr.length

---

看到有大佬的解题思路，用这个思路写一下

• 考虑逆过程，从全1数组逐个填0，直到出现合法状态时return即可
• 合法状态为第一次出现连续m个1的情况
• 每次对位置pos填0都去检查左边和右边是否有满足以下条件的段
• 1. $左区间[pos-1-m,pos-1]或右区间[pos+1,pos+1+m]$区间和为m
• 2. $左右段的边界为0$
• 注意要从后往前遍历

// #define debug
#ifdef debug
#include <time.h>
#endif

#include <iostream>
#include <algorithm>
#include <vector>
#include <string.h>
#include <map>
#include <set>
#include <stack>
#include <queue>
#include <math.h>

#define MAXN ((int)1e5+7)
#define ll long long int
#define INF (0x7f7f7f7f)
#define fori(lef, rig) for(int i=lef; i<=rig; i++)
#define forj(lef, rig) for(int j=lef; j<=rig; j++)
#define fork(lef, rig) for(int k=lef; k<=rig; k++)
#define QAQ (0)

using namespace std;

#define show(x...) \
    do { \
       cout << "\033[31;1m " << #x << " -> "; \
       err(x); \
    } while (0)

void err() { cout << "\033[39;0m" << endl; }
template<typename T, typename... A>
void err(T a, A... x) { cout << a << ' '; err(x...); }

namespace FastIO{

    char print_f[105];
    void read() {}
    void print() { putchar('\n'); }

    template <typename T, typename... T2>
       inline void read(T &x, T2 &... oth) {
           x = 0;
           char ch = getchar();
           ll f = 1;
           while (!isdigit(ch)) {
               if (ch == '-') f *= -1; 
               ch = getchar();
           }
           while (isdigit(ch)) {
               x = x * 10 + ch - 48;
               ch = getchar();
           }
           x *= f;
           read(oth...);
       }
    template <typename T, typename... T2>
       inline void print(T x, T2... oth) {
           ll p3=-1;
           if(x<0) putchar('-'), x=-x;
           do{
                print_f[++p3] = x%10 + 48;
           } while(x/=10);
           while(p3>=0) putchar(print_f[p3--]);
           putchar(' ');
           print(oth...);
       }
} // namespace FastIO
using FastIO::print;
using FastIO::read;

#define lowbit(x) (x&(-x))
int tree[MAXN], vis[MAXN], n;

void update(int i, int val) {
    while(i <= n) {
        tree[i] += val;
        i += lowbit(i);
    }
}

int query(int i) {
    int ret = 0;
    while(i > 0) {
        ret += tree[i];
        i -= lowbit(i);
    }
    return ret;
}

#define cls(x) (memset(x, 0, sizeof(x)))
//query_lr(L,R) 用来查询区间[L,R]的区间和
#define query_lr(L, R) (query(R)-query(L-1))

class Solution {
public:
    int findLatestStep(vector<int>& a, int M) {
        n = a.size();
        if(M == n) return n;
        memset(tree, false, sizeof(tree));
        memset(vis, false, sizeof(vis));
        int timer = n;
        //初始化树状数组
        for(auto x : a) update(x, 1), vis[x] = true;
     
        for(int i=n-1; i>=0; i--) { //从后往前遍历
            int now = a[i], L = a[i] - 1, R = a[i] + 1;
            vis[now] = false;
            update(now, -1); //更新当前点
            timer --;
            if(L-M+1 > 0) {
                int sum = query_lr(L-M+1, L); //查询左区间和
                if(sum == M && !vis[L-M]) return timer;
            }
            if(R+M-1 <= n) {
                int sum = query_lr(R, R+M-1); //查询右区间和
                if(sum == M && !vis[R+M]) return timer;
            }
        }
        return -1;
    }
};


#ifdef debug
signed main() {

    // vector<int> vec = { 3, 5, 1, 2, 4 };
    // vector<int> vec = { 4, 3, 2, 1 };
    vector<int> vec = {4,3,2,1};
    int M = 1;
    Solution s;
    cout << s.findLatestStep(vec, M) << endl;


   return 0;
}
#endif