java - 查询和更新 A[L]*1 + A[L+1]*2 ... + A[R]*(R-L+1) 的高效方法

Question

我们得到一个数组 A[N]，我们被要求对给定的数组执行以下 3 种类型的查询：

1个LR

2 左路

3 XY

查询1，计算： A[L] x 1 + A[L + 1] x 2 + ........ + A[R] x (R - L+ 1)

查询2，计算： A[L] x (R - L + 1) +........+ A[R - 1] x 2 + A[R] x 1

查询 3，更新： A[X] = Y

输入：

N和Q，其中N是数组中的元素数， Q是查询数。

下一行包含数组的N个元素。

接下来的Q行包含以下格式的查询：类型 LR

输出：

我们必须打印查询类型1和2的结果。

约束：

1 <= N, Q <= 10^5

1 <= L <= R <= N

1 <= A[i] <= 10^9

示例测试用例：

Inputs:

5 4
1 2 3 4 5
1 1 2
2 1 2
3 2 1
1 1 2

Output:
5
4
3

 import java.util.*;

    public class TestClass {

        public static void main(String[] args) {
            Scanner scn = new Scanner(System.in);
            int n = scn.nextInt();   
            int q = scn.nextInt();
            int[] arr = new int[n];

            for (int i = 0; i < n; i++) {
                arr[i] = scn.nextInt();
            }

            while (q-- > 0) {
                int type = scn.nextInt();    // type is for which type of query user enters.
                int l = scn.nextInt();
                int r = scn.nextInt();       
                long ans = 0;                   // ans variable i take it in long as answer may be grater.
                if (type == 1) {      
                    int count = 1;
                    while (l <= r) {
                        ans = ans + arr[l - 1] * count;  //Also in this array is started with index 1 not with 0 that means 0 index element is first element.
                        count++;
                        l++;
                    }
                    System.out.println(ans);  //Printing answer for query1.
                } else if (type == 2) {
                    int count = (r - l + 1);
                    while (l <= r) {
                        ans = ans + arr[l - 1] * count;
                        count--;
                        l++;
                    }
                    System.out.println(ans);  // Printing answer here for query 2.
                } else if (type == 3) {   
                    arr[l - 1] = r;
                }
            }
        } 
    }

代码说明：

上面的代码是计算查询类型 1和类型 2的结果的简单暴力方法，我想不出任何算法或数据结构来优化查询结果。

代码时间复杂度： O(Q x N)

根据约束，看起来我们必须在O( sqrt(N) )或O( log(N) )时间内完成查询。

我们如何优化类型 1 和类型 2 查询？

标签： javaalgorithmoptimizationbrute-force

Answer 1

假设我们有值：

0 1 2 3 4 (indexes)
a b c d e

a + 2b + 3c + 4d + 5e = S

要计算 [1, 3] 的查询 1（从零开始）：

S - a - 5e - (b + c + d)

要计算 [2, 4] 的查询 1（从零开始）：

S - (a + 2b) - (2c + 2d + 2e)

所以这些类型的查询可以通过O(1)预处理来回答。我们需要两个结构：一个存储常规前缀和：

[1, 2, 3, 4, 5]
[1, 3, 6, 10, 15]

另一个是乘数：

[1, 2, 3, 4, 5]
[1, 5, 14, 30, 55]

由于我们还需要修改前缀和，我们可以为每种类型的前缀和使用一个段树，然后O(log n)及时查询和更新。

JavaScript 代码：

// Adapted from https://codeforces.com/blog/entry/18051
function build(t, n) {  // build the tree
  for (let i = n - 1; i > 0; --i) t[i] = t[i<<1] + t[i<<1|1];
}

function modify(t, p, value, n) {  // set value at position p
  for (t[p += n] = value; p > 1; p >>= 1) t[p>>1] = t[p] + t[p^1];
}

function query(t, l, r, n) {  // sum on interval [l, r)
  let res = 0;
  for (l += n, r += n; l < r; l >>= 1, r >>= 1) {
    if (l&1) res += t[l++];
    if (r&1) res += t[--r];
  }
  return res;
}
// End adaptation from https://codeforces.com/blog/entry/18051


function main() {
  var A = [1,2,3,4,5];
  var Qs = [
    [1, 1, 2],
    [2, 1, 2],
    [3, 2, 1],
    [1, 1, 2]
  ];

  const n = A.length;
  const t0 = new Array(2 * (n + 1)).fill(0);
  const t1 = new Array(2 * (n + 1)).fill(0);
  const t2 = new Array(2 * (n + 1)).fill(0);

  // Build segment trees
  for (let i = 0; i < A.length; i++){
    t0[n + i] = A[i];
    t1[n + i] = (i + 1) * A[i];
    t2[n + i] = (n - i) * A[i];
  }
  build(t0, n);
  build(t1, n);
  build(t2, n);

  for (let [type, lx, ry] of Qs){
    // Adjust for non-zero-based indexes
    lx = lx - 1;
    ry = ry - 1;

    if (type == 1){
      let S = query(t1, 0, n + 1, n);
      let left = query(t1, 0, lx, n);
      let right = query(t1, ry + 1, n + 1, n);
      let subtrahend = lx * query(t0, lx, ry + 1, n);
      console.log(S - left - right - subtrahend);
    
    } else if (type == 2){
      let S = query(t2, 0, n + 1, n);
      let left = query(t2, 0, lx, n);
      let right = query(t2, ry + 1, n + 1, n);
      let subtrahend = (n - ry - 1) * query(t0, lx, ry + 1, n);
      console.log(S - left - right - subtrahend);

    } else {
      ry = ry + 1;
      modify(t0, lx, ry, n);
      modify(t1, lx, ry * (lx + 1), n);
      modify(t2, lx, ry * (n - lx), n);
    }
  }
}

main();