【題解】TCIRC d101: Q-7-11. 紅白彩帶 (APCS201902)

#include <iostream>
#include <set>
using namespace std;

const int maxn = 100005;
int n, k, c[maxn], p[maxn], sum_mx, sum_mn;
multiset <int> st;

void merge(int x, int y){
    int px = p[x], py = p[y];
    //移除兩段舊的長度
    st.erase(st.find(x-px+1));
    st.erase(st.find(py-y+1));
    //利用陣列把屬於同一段紅色彩帶的格子串成一個環
    p[px] = py;
    p[py] = px;
    //加入這段新長度
    st.insert(py-px+1);
}

int main() {
    cin >> n >> k;
    for (int i = 1; i <= n; i++){
        cin >> c[i];
        if (c[i] == 1){
            p[i] = i;
            st.insert(1);
            if (c[i-1] == 1) merge(i-1, i);
        }
    }
    if (st.size()){
        sum_mx += *st.rbegin();
        sum_mn += *st.begin();
    }
    for (int i = 0; i < k; i++){
        int q;
        cin >> q;
        c[q] = 1;
        p[q] = q;
        st.insert(1);
        if (c[q-1] == 1) merge(q-1, q);
        if (c[q+1] == 1) merge(q, q+1);
        sum_mx += *st.rbegin();
        sum_mn += *st.begin();
    }
    cout << sum_mx << "\n" << sum_mn << "\n";
}

#include <iostream>
#include <set>
using namespace std;
const int maxn = 100005;
int n, k;
int p[maxn], color[maxn];
multiset<int> st;

int Find(int x) {
    //查找同時進行路徑壓縮
    return p[x] < 0? x : p[x] = Find(p[x]);
}

void Union(int x, int y) {
    int g1 = Find(x);
    int g2 = Find(y);
    //把小群組併進大群組
    if (p[g1] > p[g2]) swap(g1, g2);
    //利用multiset紀錄每段紅色彩帶的長度
    auto it = st.find(-p[g1]);
    st.erase(it);
    it = st.find(-p[g2]);
    st.erase(it);
    p[g1] += p[g2];
    p[g2] = g1;
    st.insert(-p[g1]);
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> n >> k;
    for (int i = 1; i <= n; i++) {
        cin >> color[i];
        if (color[i]) {
            p[i] = -1;
            st.insert(1);
            if (color[i-1]) Union(i-1, i);
        }
    }
    //利用multiset紀錄每段紅色彩帶的長度，
    //可以快速找出最短及最長的長度
    int mi = *st.rbegin();
    int ma = *st.begin();
    int pos;
    for (int i = 0; i < k; i++) {
        cin >> pos;
        color[pos] = 1;
        p[pos] = -1;
        st.insert(1);
        if (color[pos-1]) Union(pos-1, pos);
        if (color[pos+1]) Union(pos, pos+1);
        mi += *st.rbegin();
        ma += *st.begin();
    }
    cout << mi << "\n";
    cout << ma << "\n";
    return 0;
}