백준 26306 - Balanced Seesaw Array

Updated:

Categories:

Tags:

백준 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일 때만 가능하다.

코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
#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;
}

Leave a comment