Count the number of houses at a certain distance ii

impl Solution {
  pub fn count_of_pairs(n: i32, x: i32, y: i32) -> Vec<i64> {
    let n = n as usize;
    // Normalize to 0-indexed with x <= y
    let (x, y) = {
      let (a, b) = (x as usize - 1, y as usize - 1);
      (a.min(b), a.max(b))
    };

    // Base case: pure line, result[k] = 2*(n - (k+1)) for k in 0..n
    let mut res: Vec<i64> = (1..=n).map(|k| 2 * (n - k) as i64).collect();

    let gap = y - x;
    if gap <= 1 {
      return res; // shortcut is trivial or same edge
    }

    // adj[k] accumulates net change to res[k]
    let mut adj = vec![0i64; n + 2];

    // --- Segment A: i in [0,x], j in [y,n-1] ---
    // new_dist s = (x-i) + 1 + (j-y) = a+b+1, a in [0,x], b in [0,n-1-y]
    // old_dist = gap + s - 1  (constant gain = gap-1 for all pairs in A)
    // cnt_A(s) = #{(a,b): a+b=s-1, 0<=a<=x, 0<=b<=n-1-y}
    let q = n - 1 - y; // max value of b
    let max_s = x + q + 1;
    for s in 1..=max_s {
      let a_hi = (s - 1).min(x);
      let a_lo = if s - 1 > q { s - 1 - q } else { 0 };
      if a_hi >= a_lo {
        let cnt = (a_hi - a_lo + 1) as i64;
        adj[s - 1] += 2 * cnt;
        let old_idx = gap + s - 2; // = (old_dist - 1)
        if old_idx < n {
          adj[old_idx] -= 2 * cnt;
        }
      }
    }

    // --- Segments B and C (combined via diff array) ---
    //
    // Segment B: i in [x+1, y-1], j in [y, n-1]
    //   a = i-x in [1, max_t], b = j-y in [0, q=n-1-y]
    //   new_dist s = a+b+1, old_dist = gap-a+b
    //   gain = gap-1-2a > 0 => a in [1, floor((gap-2)/2)]
    //   For fixed a: new_dist range [a+1..a+q+1], old_dist range [gap-a..gap-a+q]
    //   => diff_bc[a] += 2,     diff_bc[a+q+1] -= 2
    //      diff_bc[gap-a-1] -= 2, diff_bc[gap-a+q] += 2
    //
    // Segment C: i in [0, x], j in [x+1, y-1]
    //   c = y-j in [1, max_t], a2 = x-i in [0, x]
    //   new_dist s = a2+c+1, old_dist = gap-c+a2
    //   gain = gap-1-2c > 0 => c in [1, floor((gap-2)/2)]
    //   For fixed c: new_dist range [c+1..c+x+1], old_dist range [gap-c..gap-c+x]
    //   => diff_bc[c] += 2,     diff_bc[c+x+1] -= 2
    //      diff_bc[gap-c-1] -= 2, diff_bc[gap-c+x] += 2
    {
      let max_t = (gap - 2) / 2;
      let mut diff_bc = vec![0i64; n + 2];
      for t in 1..=max_t {
        // Segment B (a = t)
        diff_bc[t] += 2;
        if t + q + 1 <= n { diff_bc[t + q + 1] -= 2; }
        if gap - t - 1 <= n { diff_bc[gap - t - 1] -= 2; }
        if gap - t + q <= n { diff_bc[gap - t + q] += 2; }

        // Segment C (c = t)
        diff_bc[t] += 2;
        if t + x + 1 <= n { diff_bc[t + x + 1] -= 2; }
        if gap - t - 1 <= n { diff_bc[gap - t - 1] -= 2; }
        if gap - t + x <= n { diff_bc[gap - t + x] += 2; }
      }
      // Apply prefix sum of diff_bc into adj
      let mut running = 0i64;
      for k in 0..n {
        running += diff_bc[k];
        adj[k] += running;
      }
    }

    // --- Segment D: both i and j strictly between x and y ---
    //
    // i in [x+1, y-1], j in [x+1, y-1], i < j
    //   a = i-x in [1, gap-1], b = y-j in [1, gap-1], s = a+b
    //   new_dist = s+1  (go to x, shortcut, go from y to j)
    //   old_dist = gap-s
    //   gain > 0 => 2s < gap-1 => s in [2, floor((gap-2)/2)]
    //   cnt(s) = s-1  (pairs with a in [1,s-1], b=s-a in [1,s-1])
    {
      let max_s_d = (gap - 2) / 2;
      for s in 2..=max_s_d {
        let cnt = (s - 1) as i64;
        adj[s] += 2 * cnt;              // new_dist-1 = s
        let old_k = gap - s - 1;        // old_dist-1 = gap-s-1
        if old_k < n {
          adj[old_k] -= 2 * cnt;
        }
      }
    }

    // Apply adj to res
    for k in 0..n {
      res[k] += adj[k];
    }

    res
  }
}