백준 18436 수열과 쿼리 37

세그먼트 트리를 정복해보자!

그런데 문제 풀 때, even->odd / odd ->even의 경우를 고려해야 한다는 것을 늦게 알아 몇 번 틀렸다..

#include <iostream>
#include <vector>

using namespace std;

int N, M;

vector<long long> arr;
vector<int> oddTree;
vector<int> evenTree;

struct cmd {
  int op;
  int idx;
  long long x;
};

vector<cmd> query;

void input() {
  cin >> N;
  arr.resize(N, 0);
  oddTree.resize(N * 4, 0);
  evenTree.resize(N * 4, 0);

  for (int i = 0; i < N; i++) {
    cin >> arr[i];
  }

  cin >> M;
  for (int i = 0; i < M; i++) {
    int op;
    int idx;
    long long x;
    cin >> op >> idx >> x;
    query.push_back({op, idx, x});
  }
}

int init_oddTree(int start, int end, int node) {
  if (start == end) {
    if (arr[start] % 2 == 1)
      return oddTree[node] = 1;
    return oddTree[node] = 0;
  }
  int mid = (start + end) / 2;
  return oddTree[node] = init_oddTree(start, mid, node * 2) +
                         init_oddTree(mid + 1, end, node * 2 + 1);
}

int init_evenTree(int start, int end, int node) {
  if (start == end) {
    if (arr[start] % 2 == 0)
      return evenTree[node] = 1;
    return evenTree[node] = 0;
  }
  int mid = (start + end) / 2;
  return evenTree[node] = init_evenTree(start, mid, node * 2) +
                          init_evenTree(mid + 1, end, node * 2 + 1);
}

int query_evenTree(int start, int end, int node, int left, int right) {
  if (left > end || right < start)
    return 0;
  if (left <= start && right >= end)
    return evenTree[node];
  int mid = (start + end) / 2;
  return query_evenTree(start, mid, node * 2, left, right) +
         query_evenTree(mid + 1, end, node * 2 + 1, left, right);
}

int query_oddTree(int start, int end, int node, int left, int right) {
  if (left > end || right < start)
    return 0;
  if (left <= start && right >= end)
    return oddTree[node];
  int mid = (start + end) / 2;
  return query_oddTree(start, mid, node * 2, left, right) +
         query_oddTree(mid + 1, end, node * 2 + 1, left, right);
}

void update_even(int start, int end, int node, int index, int diff) {
  if (start > index || index > end)
    return;
  evenTree[node] += diff;
  if (start == end)
    return;
  int mid = (start + end) / 2;
  if (start != end) {
    update_even(start, mid, node * 2, index, diff);
    update_even(mid + 1, end, node * 2 + 1, index, diff);
  }
}

void update_odd(int start, int end, int node, int index, int diff) {
  if (start > index || index > end)
    return;
  oddTree[node] += diff;
  if (start == end)
    return;
  int mid = (start + end) / 2;
  if (start != end) {
    update_odd(start, mid, node * 2, index, diff);
    update_odd(mid + 1, end, node * 2 + 1, index, diff);
  }
}

void solve() {
  for (int i = 0; i < M; i++) {
    int op, idx;
    op = query[i].op;
    idx = query[i].idx;
    if (op == 1) {
      // A_idx를 x로 바꾸기
      long long before = arr[idx - 1];
      long long x = query[i].x;
      arr[idx - 1] = x;
      if (x % 2 == 0) {
        // even
        if (before % 2 == 1) {
          // odd였는데 even으로 바뀔 경우.
          update_even(0, N - 1, 1, idx - 1, 1);
          update_odd(0, N - 1, 1, idx - 1, -1);
        }
      } else {
        // odd
        if (before % 2 == 0) {
          // even이었는데 odd로 바뀔 경우.
          update_even(0, N - 1, 1, idx - 1, -1);
          update_odd(0, N - 1, 1, idx - 1, 1);
        }
      }
    } else if (op == 2) {
      // 짝수 개수
      int r = query[i].x;
      cout << query_evenTree(0, N - 1, 1, idx - 1, r - 1) << "\n";
    } else if (op == 3) {
      // 홀수 개수
      int r = query[i].x;
      cout << query_oddTree(0, N - 1, 1, idx - 1, r - 1) << "\n";
    }
  }
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  input();
  init_evenTree(0, N - 1, 1);
  init_oddTree(0, N - 1, 1);
  solve();
  return 0;
}