โจทย์นี้กำหนดว่าในแต่จะมีหนังสือ $N$ เล่มที่ปรากฏวันละเล่ม โดยเล่มที่ $i$ จะออกมาที่ตำแหน่ง $X_i$ และมีค่า $A_i$

โจทย์ต้องการให้เก็บมูลค่ารวมได้มากที่สุด โดยเริ่มจากตำแหน่ง $S$ ก่อนหนังสือเล่มแรก และในแต่ละวันจะเดินทางได้ระยะ $K$ แต่จะต้องเดินทางไปเก็บหนังสือในวันนั้นเท่านั้น

แนวคิด

ข้อนี้เป็นโจทย์ Segment Tree (https://programming.in.th/tasks/1147/solution)

สำหรับโจทย์ข้อนี้หากค่าของพิกัดอยู่ในช่วงที่สมเหตุสมผลจะสามารถใช้ Segment Tree ได้โดยตรง นั่นคือสำหรับหนังสือแต่ละเล่มต้อง Query ช่วงหาค่ารวมมากสุดใน $[X_i-K,X_i+K]$ และ Update ตำแหน่ง $X_i$ ให้เป็นค่าที่ได้บวก $A_i$ ซึ่งผลรวมนี้จะแทนค่าที่มากที่สุดที่เป็นไปได้หากจบที่การเก็บเล่มที่ $i$ โดยในตอนแรกมีเพียงตำแหน่ง $S$ ที่มีมูลค่ารวม 0 และที่ตำแหน่งอื่นเป็น $-\infty$

แต่เนื่องจาก $X_i$ อาจมากถึง $1,000,000,000$ จึงไม่สามารถใช้ Segment Tree โดยตรงเพราะต้องใช้ Memory $\mathcal{O}(N)$ จึงต้องใช้ Coordinate Compression

Coordinate Compression

สังเกตว่ามีเพียง $N+1$ พิกัดที่จะต้องสนใจ กล่าวคือตำแหน่ง $X_i$ ทั้งหลายและพิกัดเริ่มต้น $S$ เพราะตำแหน่งที่เหลือไม่สามารถเป็นจุดจบของการเดินทางในแต่ละวันตามที่โจทย์ระบุ

ดังนั้นเราสามารถนำพิกัดทั้งหลายมา Sort และลบค่าที่ซ้ำ และใช้ Index ใน Array ที่ได้เพื่อเป็น Index ที่ใช้ใน Segment Tree แทนพิกัด ซึ่งจะทำให้ Segment Tree เหลืออย่างมาก $N+1$ ดัชนี

ตัวอย่างการทำ Coordinate Compresssion

vector<int> coords;
coords.push_back(S);

for (int i = 0; i < N; i++) {
  cin >> X[i] >> A[i];
  coords.push_back(X[i]);
}

sort(coords.begin(), coords.end());
auto last = std::unique(coords.begin(), coords.end());
coords.erase(last, coords.end());

int C = coords.size() - 1;

unordered_map<int, int> m;
for (int i = 0; i < coords.size(); i++)
  m[coords[i]] = i;

โค้ดนี้เริ่มจากการเอาพิกัดที่ต้องการใส่ไปใน vector จากนั้น Sort และลบค่าซ้ำ สุดท้ายสร้าง unordered_map เพื่อแปลงพิกัดเดิมเป็นดัชนีใหม่

การทำ Path Compression แบบนี้จะใช้เวลา $\mathcal{O}(N \log N)$ จากการ Sort

Solution

ต่อมาเมื่อทำ Coordinate Compression แล้วสามารถใช้ตำแหน่งใหม่ในการ Update / Query ค่าใน Segment Tree ดังที่อธิบายไว้ข้างบน โดยเพียงต้องแก้ที่การ Query $[X_i-K,X_i+K]$ จะต้องแก้เป็นการหาพิกัด $l,r$ โดย $l$ เป็นดัชนีน้อยสุดที่ทำให้ $X_i-K \leq coords[l]$ และ $r$ เป็นดัชันีมากสุดที่ $coords[r] \leq X_i+K$ ซึ่งทั้งสองสามารถใช้ lower_bound กับ upper_bound ใน STL

int l = lower_bound(coords.begin(), coords.end(), X[i] - K) - coords.begin();
int r = (upper_bound(coords.begin(), coords.end(), X[i] + K) - coords.begin()) - 1;

สำหรับ Time Complexity ในการหา lower_bound กับ upper_bound ในแต่ละชั้นใช้เวลา $O(\log N)$ การ Query / Update Segment Tree ก็เช่นกัน ดังนั้นทั้งหมดจึงเป็น $O(N \log N)$

โค้ดเต็มที่เอาขั้นตอนที่อธิบายไว้มารวมกัน

#include <algorithm>
#include <vector>
#include <iostream>
#include <unordered_map>

using namespace std;

// Segment Tree
int ST[100100 * 4 + 100];

int update(int i, int Z, int n, int l, int r) {
  if (l == i && i == r) { 
    ST[n] = Z;
    return ST[n];
  }
  if (r < i || i < l) 
    return ST[n];

  int mid = (l + r) / 2;
  int new_left_value = update(i, Z, n * 2, l, mid);
  int new_right_value = update(i, Z, n * 2 + 1, mid + 1, r);

  ST[n] = max(new_left_value, new_right_value);
  return ST[n];
}

int query(int A, int B, int n, int l, int r) {
  if (A <= l && r <= B) 
    return ST[n];
  if (B < l || r < A)   
    return -1000000001; // -inf

  int mid = (l + r) / 2;
  int left_query = query(A, B, n * 2, l, mid);
  int right_query = query(A, B, n * 2 + 1, mid + 1, r);

  return max(left_query, right_query);
}

int X[100100];
int A[100100];
int M[100100];

int main() {
  int N, K, S;
  cin >> N >> K >> S;

  vector<int> coords;
  coords.push_back(S);

  for (int i = 0; i < N; i++) {
    cin >> X[i] >> A[i];
    coords.push_back(X[i]);
  }

  sort(coords.begin(), coords.end());
  auto last = std::unique(coords.begin(), coords.end());
  coords.erase(last, coords.end());

  int C = coords.size() - 1;

  unordered_map<int, int> m;
  for (int i = 0; i < coords.size(); i++)
    m[coords[i]] = i;

  for (int i = 0; i <= C; i++)
    update(i, -1000000010, 1, 0, C);

  update(m[S], 0, 1, 0, C);

  int ans = 0;

  for (int i = 0; i < N; i++) {
    int l =
        lower_bound(coords.begin(), coords.end(), X[i] - K) - coords.begin();
    int r =
        (upper_bound(coords.begin(), coords.end(), X[i] + K) - coords.begin()) -
        1;

    int M = query(l, r, 1, 0, C) + A[i];
    ans = max(ans, M);

    update(m[X[i]], M, 1, 0, C);
  }

  cout << ans;
}