Problem p02630 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <algorithm>
#include <bits/stdc++.h>
#include <unordered_set>

#define ll long long

using namespace std;

ll f(ll x) {

  for (int i = 2; i * i <= x; i++) {
    if (x % i == 0) {
      return i;
    }
  }

  return x;
}

int main() {

  ios_base::sync_with_stdio(false);
  cin.tie(NULL);

  int t;

  ll n, k;

  ll Q, b, c;

  cin >> n;

  ll sum = 0;
  vector<ll> a(n);
  map<ll, int> occur;
  for (ll i = 0; i < n; i++) {
    cin >> a[i];
    occur[a[i]]++;
    sum += a[i];
  }

  cin >> Q;

  while (Q--) {
    cin >> b >> c;

    sum += (c - b) * occur[b];
    cout << sum << endl;
    occur[c] += occur[b];
    occur[b] = 0;
  }
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <bits/stdc++.h>
#include <unordered_set>

#define ll long long

using namespace std;

ll getMid(int s, int e) { return s + (e - s) / 2; }

/* A recursive function to get the sum of values in the given range
    of the array. The following are parameters for this function.

    st --> Pointer to segment tree
    si --> Index of current node in the segment tree. Initially
            0 is passed as root is always at index 0
    ss & se --> Starting and ending indexes of the segment represented
                by current node, i.e., st[si]
    qs & qe --> Starting and ending indexes of query range */
ll getSumUtil(ll *st, ll ss, ll se, ll qs, ll qe, ll si) {
  // If segment of this node is a part of given range, then return
  // the sum of the segment
  if (qs <= ss && qe >= se)
    return st[si];

  // If segment of this node is outside the given range
  if (se < qs || ss > qe)
    return 0;

  // If a part of this segment overlaps with the given range
  ll mid = getMid(ss, se);
  return getSumUtil(st, ss, mid, qs, qe, 2 * si + 1) +
         getSumUtil(st, mid + 1, se, qs, qe, 2 * si + 2);
}

/* A recursive function to update the nodes which have the given
index in their range. The following are parameters
    st, si, ss and se are same as getSumUtil()
    i --> index of the element to be updated. This index is
            in the input array.
diff --> Value to be added to all nodes which have i in range */
void updateValueUtil(ll *st, ll ss, ll se, ll i, ll diff, ll si) {
  // Base Case: If the input index lies outside the range of
  // this segment
  if (i < ss || i > se)
    return;

  // If the input index is in range of this node, then update
  // the value of the node and its children
  st[si] = st[si] + diff;
  if (se != ss) {
    ll mid = getMid(ss, se);
    updateValueUtil(st, ss, mid, i, diff, 2 * si + 1);
    updateValueUtil(st, mid + 1, se, i, diff, 2 * si + 2);
  }
}

// The function to update a value in input array and segment tree.
// It uses updateValueUtil() to update the value in segment tree
void updateValue(vector<ll> &arr, ll *st, ll n, ll i, ll new_val) {
  // Check for erroneous input index
  if (i < 0 || i > n - 1) {
    cout << "Invalid Input";
    return;
  }

  // Get the difference between new value and old value
  ll diff = new_val - arr[i];

  // Update the value in array
  arr[i] = new_val;

  // Update the values of nodes in segment tree
  updateValueUtil(st, 0, n - 1, i, diff, 0);
}

// Return sum of elements in range from index qs (quey start)
// to qe (query end). It mainly uses getSumUtil()
ll getSum(ll *st, ll n, ll qs, ll qe) {
  // Check for erroneous input values
  if (qs < 0 || qe > n - 1 || qs > qe) {
    cout << "Invalid Input";
    return -1;
  }

  return getSumUtil(st, 0, n - 1, qs, qe, 0);
}

// A recursive function that constructs Segment Tree for array[ss..se].
// si is index of current node in segment tree st
ll constructSTUtil(vector<ll> &arr, ll ss, ll se, ll *st, ll si) {
  // If there is one element in array, store it in current node of
  // segment tree and return
  if (ss == se) {
    st[si] = arr[ss];
    return arr[ss];
  }

  // If there are more than one elements, then recur for left and
  // right subtrees and store the sum of values in this node
  int mid = getMid(ss, se);
  st[si] = constructSTUtil(arr, ss, mid, st, si * 2 + 1) +
           constructSTUtil(arr, mid + 1, se, st, si * 2 + 2);
  return st[si];
}

/* Function to construct segment tree from given array. This function
allocates memory for segment tree and calls constructSTUtil() to
fill the allocated memory */
ll *constructST(vector<ll> &arr, ll n) {
  // Allocate memory for the segment tree

  // Height of segment tree
  ll x = (ll)(ceil(log2(n)));

  // Maximum size of segment tree
  ll max_size = 2 * (ll)pow(2, x) - 1;

  ll *st = new ll[max_size];
  constructSTUtil(arr, 0, n - 1, st, 0);
  return st;
}

int main() {

  ios_base::sync_with_stdio(false);
  cin.tie(NULL);

  int n, q;

  cin >> n;
  vector<ll> a(n);
  ll sum = 0;
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }

  ll *st = constructST(a, n);

  cin >> q;

  ll b, c;
  while (q--) {
    cin >> b >> c;

    for (int i = 0; i < n; i++) {
      if (a[i] == b) {
        updateValue(a, st, n, i, c);
      }
    }

    cout << getSum(st, n, 0, n - 1) << endl;
  }

  return 0;
}