백준 26306 - Balanced Seesaw Array

백준 26306 - Balanced Seesaw Array

풀이

문제에서 주어진 공식을 정리하면 다음과 같다.

\[\Sigma_{i=1}^{m}(i\times a_i) - k \times \Sigma_{i=1}^{m}(a_i) = 0\] \[\Sigma_{i=1}^{m}(i\times a_i) = k \times \Sigma_{i=1}^{m}(a_i)\]

위를 만족시키는 k를 찾으면 해결할 수 있다.

구간 쿼리가 들어오는 걸로 보았을 때 lazy segment tree를 사용하면 된다는 것을 알 수 있다.
즉, 우리는 $\Sigma (i \times a_i)$와 $\Sigma(a_i)$를 계산하고 있는 2개의 lazy 세그를 구하면 된다.

lazy 구조체를 다음과 같이 만든다.

• val
• a
• b

lazy를 update 시키면 $a \times val + b$ 를 의미한다.
1번 쿼리가 들어오면 a=1, b=x로 업데이트 해주면 되고, 2번 쿼리가 들어오면 a=0, b=x으로 업데이트 해주면 된다.

이때 $\Sigma(a_i)$의 경우는 b가 더해질 때 해당 (start ~ end)에 대해 각각 b씩 더해지므로 $(end-start+1)\times b$를 해주면 되고,
$\Sigma (i \times a_i)$의 경우는 해당 (start ~ end)에 대해 $start\times b, (start+1)\times b, … , end\times b$ 가 더해지므로 결국, $b\times (start + (start+1) + … + end)$ 가 더해진다. 따라서 $b\times (end\times (end+1)/2 - (start-1)\times start/2)$를 해주면 된다.

a의 경우는 그냥 전체에 곱해지는 것이므로 $\Sigma (i \times a_i)$와 $\Sigma(a_i)$ 모두 그냥 곱해주면 된다.

위와 같이 세그를 구현해주면 해결할 수 있다.

마지막으로 주의할 점이 있는데,
처음의 공식의 경우 i가 1 ~ m 의 범위이고, 쿼리는 l ~ r의 범위이다.
공식을 l과 r로 변경하면 다음과 같다.

\[\Sigma_{i=1}^{r-l+1}(i*a_{i+l-1}) - k \times \Sigma_{i=l}^r(a_i) = 0\]

이를 우리의 세그에 좀 더 알맞게 정리하면 다음과 같다.

\[\Sigma_{i=l}^r(i \times a_i) - (l-1)\times \Sigma_{i=l}^r(a_i) - k \times \Sigma_{i=l}^r(a_i) = 0\]

\(\Sigma_{i=l}^r(i \times a_i) = (k+l-1)\times \Sigma_{i=l}^r(a_i)\) 여기서 k는 1 ~ m(=r-l+1) 범위 이어야 하므로 k+l-1은 l ~ r 의 범위이라면 balanced seesaw array임을 알 수 있다.

만약 $\Sigma_{i=l}^r(a_i)$가 0이라면 $\Sigma_{i=l}^r(i \times a_i)$가 0일 때만 가능하다.

코드

#include <bits/stdc++.h>

#define endl "\n"
#define all(v) (v).begin(), (v).end()
#define For(i, n) for(int i=0; i<n; ++i)
#define For1(i, n) for(int i=1; i<=n; ++i)
#define For2(i, a, b) for(int i=(a); i<=(b); ++i)
#define ft first
#define sd second
#define Get(i, v) get<i>(v)

using namespace std;
using ll = long long;
using ull = unsigned long long;
using lll = __int128;
using ulll = __uint128_t;
using ld = long double;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
using ti3 = tuple<int, int, int>;
using tl3 = tuple<ll, ll, ll>;

const int INF = numeric_limits<int>::max();
const ll LNF = numeric_limits<ll>::max();

struct Lazy {
    ll val, a, b; // a * val + b
};

class Lazy_Segment {
public:
    vector<Lazy> tree; //tree[node] := a[start ~ end] 의 합
    int TYPE;

    Lazy_Segment() {}
    Lazy_Segment(int size) {
        this->resize(size);
    }
    void resize(int size) {
        size = (int) floor(log2(size)) + 2;
        size = pow(2, size);
        tree.resize(size, {0,1,0});
    }
    ll type(ll s, ll e) const {
        if(TYPE == 1) return e-s+1;
        return e*(e+1)/2 - (s-1)*s/2;
    }
    ll init(vector<ll> &a, int node, int start, int end) {
        if(start == end) return tree[node].val = type(start,end)*a[start];
        return tree[node].val = (init(a, 2*node, start, (start+end)/2) + init(a, 2*node+1, (start+end)/2+1, end));
    }
    void update_lazy(int node, int start, int end) {
        if(tree[node].a == 1 && tree[node].b == 0) return;
        tree[node].val = (tree[node].a*tree[node].val + tree[node].b*type(start,end));
        if(start != end) {
            for(auto i : {2*node, 2*node+1}) {
                tree[i].a = (tree[node].a * tree[i].a);
                tree[i].b = (tree[node].a * tree[i].b + tree[node].b);
            }
        }
        tree[node].a = 1, tree[node].b = 0;
    }
    void update(int node, int start, int end, int left, int right, ll a, ll b) {
        update_lazy(node, start, end);
        if(right < start || end < left) return;
        if(left <= start && end <= right) {
            tree[node].a = (tree[node].a * a);
            tree[node].b = (tree[node].b + b);
            update_lazy(node, start, end);
            return;
        }
        update(node * 2, start, (start + end) / 2, left, right, a, b);
        update(node * 2 + 1, (start + end) / 2 + 1, end, left, right, a, b);
        tree[node].val = (tree[2*node].val + tree[2*node+1].val);
    }
    ll sum(int node, int start, int end, int left, int right) {
        update_lazy(node, start, end);
        if(right < start || end < left) return 0;
        if(left <= start && end <= right) return tree[node].val;
        return (sum(node * 2, start, (start + end) / 2, left, right) +
                sum(node * 2 + 1, (start + end) / 2 + 1, end, left, right));
    }
};

void solve() {
    int n, q; cin >> n >> q;
    vector<ll> a(n+1);
    For1(i,n) cin >> a[i];
    Lazy_Segment seg1(n), seg2(n);
    seg1.TYPE = 1, seg2.TYPE = 2;
    seg1.init(a,1,1,n); seg2.init(a,1,1,n);
    while(q--) {
        int cmd, l, r; cin >> cmd >> l >> r;
        if(cmd == 1) {
            ll x; cin >> x;
            seg1.update(1,1,n,l,r,1,x);
            seg2.update(1,1,n,l,r,1,x);
        }
        else if(cmd == 2) {
            ll x; cin >> x;
            seg1.update(1,1,n,l,r,0,x);
            seg2.update(1,1,n,l,r,0,x);
        }
        else {
            ll res1 = seg1.sum(1,1,n,l,r);
            ll res2 = seg2.sum(1,1,n,l,r);
            bool flag = false;
            if(res1 == 0) {
                if(res2 == 0) flag = true;
            }
            else {
                ll k = res2/res1;
                if(l<=k and k<=r and res2 == k*res1) flag = true;
            }
            if(flag) cout << "Yes\n";
            else cout << "No\n";
        }
    }
}

int main(void) {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int tc = 1;
//    cin >> tc;
    while(tc--) {
        solve();
//        cout << solve() << endl;
    }


    return 0;
}