✨ feat(binary-search-tree): export createBinaryIndexTree1Mod and createBinaryIndexTree2Mod

Author: guanghechen

Diff for: packages/binary-index-tree/README-zh.md

@@ -204,6 +204,28 @@
   bit.query(/* any integer between [5, 10] */) // => 0n
   ```
 
+* With Mod
+
+  ```typescript
+  import { createBinaryIndexTree1Mod } from '@algorithm.ts/binary-index-tree'
+
+  const MOD = 1e9 + 7
+  const bit = createBinaryIndexTree1Mod<number>(0, MOD) 
+
+  bit.add(2, <value>)   // <value> should in the range of (-MOD, MOD)
+  bit.query(3)
+  ```
+
+  ```typescript
+  import { createBinaryIndexTree2Mod } from '@algorithm.ts/binary-index-tree'
+
+  const MOD = 1e9 + 7
+  const bit = createBinaryIndexTree1Mod<bigint>(0, BigInt(MOD)) 
+
+  bit.add(2, <value>)   // <value> should in the range of (-MOD, MOD)
+  bit.query(3)
+  ```
+
 
 ## Related
 

Diff for: packages/binary-index-tree/README.md

@@ -211,6 +211,28 @@ implementation is simpler and easier to understand.
   bit.query(/* any integer between [5, 10] */) // => 0n
   ```
 
+* With Mod
+
+  ```typescript
+  import { createBinaryIndexTree1Mod } from '@algorithm.ts/binary-index-tree'
+
+  const MOD = 1e9 + 7
+  const bit = createBinaryIndexTree1Mod<number>(0, MOD) 
+
+  bit.add(2, <value>)   // <value> should in the range of (-MOD, MOD)
+  bit.query(3)
+  ```
+
+  ```typescript
+  import { createBinaryIndexTree2Mod } from '@algorithm.ts/binary-index-tree'
+
+  const MOD = 1e9 + 7
+  const bit = createBinaryIndexTree1Mod<bigint>(0, BigInt(MOD)) 
+
+  bit.add(2, <value>)   // <value> should in the range of (-MOD, MOD)
+  bit.query(3)
+  ```
+
 
 ## Related
 

Diff for: packages/binary-index-tree/__test__/binary-saerch-tree.spec.ts

@@ -1,5 +1,11 @@
 import { randomInt } from 'crypto'
-import { createBinaryIndexTree1, createBinaryIndexTree2, lowbit } from '../src'
+import {
+  createBinaryIndexTree1,
+  createBinaryIndexTree1Mod,
+  createBinaryIndexTree2,
+  createBinaryIndexTree2Mod,
+  lowbit,
+} from '../src'
 
 describe('tree1', function () {
   const MAX_N = 1000
@@ -143,6 +149,150 @@ describe('tree2', function () {
   })
 })
 
+describe('tree1-mod', function () {
+  const MAX_N = 1000
+  const MOD = 1000 + 17
+
+  test('number', function () {
+    const bit = createBinaryIndexTree1Mod<number>(0, MOD)
+    bit.init(MAX_N)
+
+    const A: number[] = new Array(MAX_N)
+      .fill(0)
+      .map((_x, i) => randomInt(MOD - 1))
+    for (let x = 0; x < MAX_N; ++x) bit.add(x + 1, A[x])
+
+    const getSum = (x: number): number => {
+      let result = 0
+      for (let i = 0; i < x; ++i) result += A[i]
+      return result % MOD
+    }
+
+    for (let q = 0; q < 1000; ++q) {
+      const x = Math.max(1, Math.ceil(Math.random() * MAX_N))
+      const value = Math.round(Math.random() * (MOD - 1))
+
+      A[x - 1] += value
+      bit.add(x, value)
+      expect(bit.query(x)).toBe(getSum(x))
+    }
+
+    // Test illegal input.
+    bit.add(0, 1)
+    bit.add(-1, 1)
+    bit.query(0)
+    bit.query(-1)
+  })
+
+  test('bigint', function () {
+    const bit = createBinaryIndexTree1Mod<bigint>(0n, BigInt(MOD))
+    bit.init(MAX_N)
+
+    const A: Array<bigint> = new Array(MAX_N)
+      .fill(0)
+      .map((_x, i) => BigInt(randomInt(MOD)))
+    for (let x = 0; x < MAX_N; ++x) bit.add(x + 1, A[x])
+
+    const getSum = (x: number): bigint => {
+      let result = 0n
+      for (let i = 0; i < x; ++i) result += A[i]
+      return result % BigInt(MOD)
+    }
+
+    for (let q = 0; q < 1000; ++q) {
+      const x = Math.max(1, Math.ceil(Math.random() * MAX_N))
+      const value = BigInt(Math.round(Math.random() * (MOD - 1)))
+
+      A[x - 1] += value
+      bit.add(x, value)
+      expect(bit.query(x)).toBe(getSum(x))
+    }
+
+    // Test illegal input.
+    bit.add(0, 1n)
+    bit.add(-1, 1n)
+    bit.query(0)
+    bit.query(-1)
+  })
+})
+
+describe('tree2-mod', function () {
+  const MAX_N = 1000
+  const MOD = 1000 + 17
+
+  test('number', function () {
+    const bit = createBinaryIndexTree2Mod<number>(0, MOD)
+    bit.init(MAX_N)
+
+    // Initialize
+    const A: number[] = new Array(MAX_N)
+      .fill(0)
+      .map((_x, i) => randomInt(MOD - 1))
+    for (let x = 0; x < MAX_N; ++x) {
+      bit.add(x, -A[x])
+      bit.add(x + 1, A[x])
+    }
+
+    // Test for initialization.
+    for (let x = 1; x <= MAX_N; ++x) expect(bit.query(x)).toBe(A[x - 1])
+
+    const performAdd = (x: number, value: number): void => {
+      for (let i = 0; i < x; ++i) A[i] += value
+    }
+
+    for (let q = 0; q < 1000; ++q) {
+      const x = Math.max(1, Math.ceil(Math.random() * MAX_N))
+      const value = Math.round(Math.random() * (MOD - 1))
+
+      performAdd(x, value)
+      bit.add(x, value)
+      expect(bit.query(x)).toBe(A[x - 1] % MOD)
+    }
+
+    // Test illegal input.
+    bit.add(0, 1)
+    bit.add(-1, 1)
+    bit.query(0)
+    bit.query(1)
+  })
+
+  test('bigint', function () {
+    const bit = createBinaryIndexTree2Mod<bigint>(0n, BigInt(MOD))
+    bit.init(MAX_N)
+
+    // Initialize
+    const A: Array<bigint> = new Array(MAX_N)
+      .fill(0)
+      .map((_x, i) => BigInt(randomInt(MOD - 1)))
+    for (let x = 0; x < MAX_N; ++x) {
+      bit.add(x, -A[x])
+      bit.add(x + 1, A[x])
+    }
+
+    // Test for initialization.
+    for (let x = 1; x <= MAX_N; ++x) expect(bit.query(x)).toBe(A[x - 1])
+
+    const performAdd = (x: number, value: bigint): void => {
+      for (let i = 0; i < x; ++i) A[i] += value
+    }
+
+    for (let q = 0; q < 1000; ++q) {
+      const x = Math.max(1, Math.ceil(Math.random() * MAX_N))
+      const value = BigInt(Math.round(Math.random() * (MOD - 1)))
+
+      performAdd(x, value)
+      bit.add(x, value)
+      expect(bit.query(x)).toBe(A[x - 1] % BigInt(MOD))
+    }
+
+    // Test illegal input.
+    bit.add(0, 1n)
+    bit.add(-1, 1n)
+    bit.query(0)
+    bit.query(-1)
+  })
+})
+
 describe('util', function () {
   test('lowbit', function () {
     const nums: number[] = new Array(1000).fill(0).map((_x, i) => i + 1)

Diff for: packages/binary-index-tree/src/index.ts

@@ -1,4 +1,6 @@
 export * from './tree1'
+export * from './tree1-mod'
 export * from './tree2'
+export * from './tree2-mod'
 export * from './types'
 export * from './util'

Diff for: packages/binary-index-tree/src/tree1-mod.ts

@@ -0,0 +1,69 @@
+import type { BinaryIndexTree } from './types'
+import { lowbit } from './util'
+
+/**
+ * 创建一棵支持 *单点更新，区间查询* 的树状数组，所有的更新和查询操作均需对某个
+ * 数进行取模，且所有数值应在范围 (-MOD, MOD) 之间。
+ *
+ * Create a binary search tree for *single-point update with interval query*.
+ * All update and query operations need to be modulo a certain number, and all
+ * values should be in the range (-MOD, MOD).
+ *
+ * @param ZERO    0 for number, 0n for bigint.
+ * @param MOD
+ * @returns
+ */
+export function createBinaryIndexTree1Mod<T extends number | bigint>(
+  ZERO: T,
+  MOD: T,
+): BinaryIndexTree<T> {
+  let _size = 0
+  const _nodes: T[] = [ZERO]
+  return { init, add, query }
+
+  /**
+   * Initialize the BinarySearchTree.
+   * @param N
+   */
+  function init(N: number): void {
+    _size = N
+    if (_nodes.length <= _size) _nodes.length = _size + 1
+    _nodes.fill(ZERO, 1, _size + 1)
+  }
+
+  /**
+   * Add value to the xth number
+   * @param xth
+   * @param value   value \in (-MOD, MOD)
+   */
+  function add(xth: number, value: T): void {
+    if (xth < 1) return
+    // eslint-disable-next-line no-param-reassign
+    if (value < 0) value += MOD as any
+    for (let i = xth; i <= _size; i += lowbit(i)) {
+      _nodes[i] = modAdd(_nodes[i], value)
+    }
+  }
+
+  /**
+   * Calculate the prefix sum of the first x elements of the array.
+   * @param xth
+   */
+  function query(xth: number): T {
+    let result = ZERO
+    for (let i = xth; i > 0; i -= lowbit(i)) {
+      result = modAdd(result, _nodes[i])
+    }
+    return result
+  }
+
+  /**
+   * @param x   x \in [0, MOD)
+   * @param y   y \in [0, MOD)
+   * @returns
+   */
+  function modAdd(x: T, y: T): T {
+    const result: any = (x as any) + (y as any)
+    return result >= MOD ? result - MOD : result
+  }
+}

Diff for: packages/binary-index-tree/src/tree1.ts

@@ -6,7 +6,6 @@ import { lowbit } from './util'
  *
  * Create a binary search tree for *single-point update with interval query*.
  *
- * @param MAX_N
  * @param ZERO    0 for number, 0n for bigint.
  * @returns
  */
@@ -15,6 +14,7 @@ export function createBinaryIndexTree1<T extends number | bigint>(
 ): BinaryIndexTree<T> {
   let _size = 0
   const _nodes: T[] = [ZERO]
+  return { init, add, query }
 
   /**
    * Initialize the BinarySearchTree.
@@ -28,27 +28,25 @@ export function createBinaryIndexTree1<T extends number | bigint>(
 
   /**
    * Add value to the xth number
-   * @param x
+   * @param xth
    * @param value
    */
-  function add(x: number, value: T): void {
-    if (x < 1) return
-    for (let i = x; i <= _size; i += lowbit(i)) {
+  function add(xth: number, value: T): void {
+    if (xth < 1) return
+    for (let i = xth; i <= _size; i += lowbit(i)) {
       _nodes[i] += value as any
     }
   }
 
   /**
    * Calculate the prefix sum of the first x elements of the array.
-   * @param x
+   * @param xth
    */
-  function query(x: number): T {
+  function query(xth: number): T {
     let result = ZERO
-    for (let i = x; i > 0; i -= lowbit(i)) {
+    for (let i = xth; i > 0; i -= lowbit(i)) {
       result += _nodes[i] as any
     }
     return result
   }
-
-  return { init, add, query }
 }