返回介绍

solution / 2600-2699 / 2612.Minimum Reverse Operations / README_EN

发布于 2024-06-17 01:03:02 字数 45188 浏览 0 评论 0 收藏 0

2612. Minimum Reverse Operations

中文文档

Description

You are given an integer n and an integer p in the range [0, n - 1]. Representing a 0-indexed array arr of length n where all positions are set to 0's, except position p which is set to 1.

You are also given an integer array banned containing some positions from the array. For the ith position in banned, arr[banned[i]] = 0, and banned[i] != p.

You can perform multiple operations on arr. In an operation, you can choose a subarray with size k and reverse the subarray. However, the 1 in arr should never go to any of the positions in banned. In other words, after each operation arr[banned[i]] remains 0.

_Return an array_ ans _where__ for each _i_ from _[0, n - 1], ans[i] _is the minimum number of reverse operations needed to bring the_ 1 _to position_ i_ in arr_, _or_ -1 _if it is impossible_.

  • A subarray is a contiguous non-empty sequence of elements within an array.
  • The values of ans[i] are independent for all i's.
  • The reverse of an array is an array containing the values in reverse order.

 

Example 1:

Input: n = 4, p = 0, banned = [1,2], k = 4
Output: [0,-1,-1,1]
Explanation: In this case k = 4 so there is only one possible reverse operation we can perform, which is reversing the whole array. Initially, 1 is placed at position 0 so the amount of operations we need for position 0 is 0. We can never place a 1 on the banned positions, so the answer for positions 1 and 2 is -1. Finally, with one reverse operation we can bring the 1 to index 3, so the answer for position 3 is 1. 

Example 2:

Input: n = 5, p = 0, banned = [2,4], k = 3
Output: [0,-1,-1,-1,-1]
Explanation: In this case the 1 is initially at position 0, so the answer for that position is 0. We can perform reverse operations of size 3. The 1 is currently located at position 0, so we need to reverse the subarray [0, 2] for it to leave that position, but reversing that subarray makes position 2 have a 1, which shouldn't happen. So, we can't move the 1 from position 0, making the result for all the other positions -1. 

Example 3:

Input: n = 4, p = 2, banned = [0,1,3], k = 1
Output: [-1,-1,0,-1]
Explanation: In this case we can only perform reverse operations of size 1. So the 1 never changes its position.

 

Constraints:

  • 1 <= n <= 105
  • 0 <= p <= n - 1
  • 0 <= banned.length <= n - 1
  • 0 <= banned[i] <= n - 1
  • 1 <= k <= n 
  • banned[i] != p
  • all values in banned are unique 

Solutions

Solution 1: Ordered Set + BFS

We notice that for any index $i$ in the subarray interval $[l,..r]$, the flipped index $j = l + r - i$.

If the subarray moves one position to the right, then $j = l + 1 + r + 1 - i = l + r - i + 2$, that is, $j$ will increase by $2$.

Similarly, if the subarray moves one position to the left, then $j = l - 1 + r - 1 - i = l + r - i - 2$, that is, $j$ will decrease by $2$.

Therefore, for a specific index $i$, all its flipped indices form an arithmetic progression with common difference $2$, that is, all the flipped indices have the same parity.

Next, we consider the range of values ​​of the index $i$ after flipping $j$.

  • If the boundary is not considered, the range of values ​​of $j$ is $[i - k + 1, i + k - 1]$.
  • If the subarray is on the left, then $[l, r] = [0, k - 1]$, so the flipped index $j$ of $i$ is $0 + k - 1 - i$, that is, $j = k - i - 1$, so the left boundary $mi = max(i - k + 1, k - i - 1)$.
  • If the subarray is on the right, then $[l, r] = [n - k, n - 1]$, so the flipped index $j= n - k + n - 1 - i$ is $j = n \times 2 - k - i - 1$, so the right boundary of $j$ is $mx = min(i + k - 1, n \times 2 - k - i - 1)$.

We use two ordered sets to store all the odd indices and even indices to be searched, here we need to exclude the indices in the array $banned$ and the index $p$.

Then we use BFS to search, each time searching all the flipped indices $j$ of the current index $i$, that is, $j = mi, mi + 2, mi + 4, \dots, mx$, updating the answer of index $j$ and adding index $j$ to the search queue, and removing index $j$ from the corresponding ordered set.

When the search is over, the answer to all indices can be obtained.

The time complexity is $O(n \times \log n)$ and the space complexity is $O(n)$. Where $n$ is the given array length in the problem.

from sortedcontainers import SortedSet


class Solution:
  def minReverseOperations(
    self, n: int, p: int, banned: List[int], k: int
  ) -> List[int]:
    ans = [-1] * n
    ans[p] = 0
    ts = [SortedSet() for _ in range(2)]
    for i in range(n):
      ts[i % 2].add(i)
    ts[p % 2].remove(p)
    for i in banned:
      ts[i % 2].remove(i)
    ts[0].add(n)
    ts[1].add(n)
    q = deque([p])
    while q:
      i = q.popleft()
      mi = max(i - k + 1, k - i - 1)
      mx = min(i + k - 1, n * 2 - k - i - 1)
      s = ts[mi % 2]
      j = s.bisect_left(mi)
      while s[j] <= mx:
        q.append(s[j])
        ans[s[j]] = ans[i] + 1
        s.remove(s[j])
        j = s.bisect_left(mi)
    return ans
class Solution {
  public int[] minReverseOperations(int n, int p, int[] banned, int k) {
    int[] ans = new int[n];
    TreeSet<Integer>[] ts = new TreeSet[] {new TreeSet<>(), new TreeSet<>()};
    for (int i = 0; i < n; ++i) {
      ts[i % 2].add(i);
      ans[i] = i == p ? 0 : -1;
    }
    ts[p % 2].remove(p);
    for (int i : banned) {
      ts[i % 2].remove(i);
    }
    ts[0].add(n);
    ts[1].add(n);
    Deque<Integer> q = new ArrayDeque<>();
    q.offer(p);
    while (!q.isEmpty()) {
      int i = q.poll();
      int mi = Math.max(i - k + 1, k - i - 1);
      int mx = Math.min(i + k - 1, n * 2 - k - i - 1);
      var s = ts[mi % 2];
      for (int j = s.ceiling(mi); j <= mx; j = s.ceiling(mi)) {
        q.offer(j);
        ans[j] = ans[i] + 1;
        s.remove(j);
      }
    }
    return ans;
  }
}
class Solution {
public:
  vector<int> minReverseOperations(int n, int p, vector<int>& banned, int k) {
    vector<int> ans(n, -1);
    ans[p] = 0;
    set<int> ts[2];
    for (int i = 0; i < n; ++i) {
      ts[i % 2].insert(i);
    }
    ts[p % 2].erase(p);
    for (int i : banned) {
      ts[i % 2].erase(i);
    }
    ts[0].insert(n);
    ts[1].insert(n);
    queue<int> q{{p}};
    while (!q.empty()) {
      int i = q.front();
      q.pop();
      int mi = max(i - k + 1, k - i - 1);
      int mx = min(i + k - 1, n * 2 - k - i - 1);
      auto& s = ts[mi % 2];
      auto it = s.lower_bound(mi);
      while (*it <= mx) {
        int j = *it;
        ans[j] = ans[i] + 1;
        q.push(j);
        it = s.erase(it);
      }
    }
    return ans;
  }
};
func minReverseOperations(n int, p int, banned []int, k int) []int {
  ans := make([]int, n)
  ts := [2]*redblacktree.Tree{redblacktree.NewWithIntComparator(), redblacktree.NewWithIntComparator()}
  for i := 0; i < n; i++ {
    ts[i%2].Put(i, struct{}{})
    ans[i] = -1
  }
  ans[p] = 0
  ts[p%2].Remove(p)
  for _, i := range banned {
    ts[i%2].Remove(i)
  }
  ts[0].Put(n, struct{}{})
  ts[1].Put(n, struct{}{})
  q := []int{p}
  for len(q) > 0 {
    i := q[0]
    q = q[1:]
    mi := max(i-k+1, k-i-1)
    mx := min(i+k-1, n*2-k-i-1)
    s := ts[mi%2]
    for x, _ := s.Ceiling(mi); x.Key.(int) <= mx; x, _ = s.Ceiling(mi) {
      j := x.Key.(int)
      s.Remove(j)
      ans[j] = ans[i] + 1
      q = append(q, j)
    }
  }
  return ans
}
function minReverseOperations(n: number, p: number, banned: number[], k: number): number[] {
  const ans = new Array(n).fill(-1);
  const ts = new Array(2).fill(0).map(() => new TreeSet<number>());
  for (let i = 0; i < n; ++i) {
    ts[i % 2].add(i);
  }
  ans[p] = 0;
  ts[p % 2].delete(p);
  for (const i of banned) {
    ts[i % 2].delete(i);
  }
  ts[0].add(n);
  ts[1].add(n);
  let q = [p];
  while (q.length) {
    const t: number[] = [];
    for (const i of q) {
      const mi = Math.max(i - k + 1, k - i - 1);
      const mx = Math.min(i + k - 1, n * 2 - k - i - 1);
      const s = ts[mi % 2];
      for (let j = s.ceil(mi)!; j <= mx; j = s.ceil(j)!) {
        t.push(j);
        ans[j] = ans[i] + 1;
        s.delete(j);
      }
    }
    q = t;
  }
  return ans;
}

type Compare<T> = (lhs: T, rhs: T) => number;

class RBTreeNode<T = number> {
  data: T;
  count: number;
  left: RBTreeNode<T> | null;
  right: RBTreeNode<T> | null;
  parent: RBTreeNode<T> | null;
  color: number;
  constructor(data: T) {
    this.data = data;
    this.left = this.right = this.parent = null;
    this.color = 0;
    this.count = 1;
  }

  sibling(): RBTreeNode<T> | null {
    if (!this.parent) return null; // sibling null if no parent
    return this.isOnLeft() ? this.parent.right : this.parent.left;
  }

  isOnLeft(): boolean {
    return this === this.parent!.left;
  }

  hasRedChild(): boolean {
    return (
      Boolean(this.left && this.left.color === 0) ||
      Boolean(this.right && this.right.color === 0)
    );
  }
}

class RBTree<T> {
  root: RBTreeNode<T> | null;
  lt: (l: T, r: T) => boolean;
  constructor(compare: Compare<T> = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0)) {
    this.root = null;
    this.lt = (l: T, r: T) => compare(l, r) < 0;
  }

  rotateLeft(pt: RBTreeNode<T>): void {
    const right = pt.right!;
    pt.right = right.left;

    if (pt.right) pt.right.parent = pt;
    right.parent = pt.parent;

    if (!pt.parent) this.root = right;
    else if (pt === pt.parent.left) pt.parent.left = right;
    else pt.parent.right = right;

    right.left = pt;
    pt.parent = right;
  }

  rotateRight(pt: RBTreeNode<T>): void {
    const left = pt.left!;
    pt.left = left.right;

    if (pt.left) pt.left.parent = pt;
    left.parent = pt.parent;

    if (!pt.parent) this.root = left;
    else if (pt === pt.parent.left) pt.parent.left = left;
    else pt.parent.right = left;

    left.right = pt;
    pt.parent = left;
  }

  swapColor(p1: RBTreeNode<T>, p2: RBTreeNode<T>): void {
    const tmp = p1.color;
    p1.color = p2.color;
    p2.color = tmp;
  }

  swapData(p1: RBTreeNode<T>, p2: RBTreeNode<T>): void {
    const tmp = p1.data;
    p1.data = p2.data;
    p2.data = tmp;
  }

  fixAfterInsert(pt: RBTreeNode<T>): void {
    let parent = null;
    let grandParent = null;

    while (pt !== this.root && pt.color !== 1 && pt.parent?.color === 0) {
      parent = pt.parent;
      grandParent = pt.parent.parent;

      /*  Case : A
        Parent of pt is left child of Grand-parent of pt */
      if (parent === grandParent?.left) {
        const uncle = grandParent.right;

        /* Case : 1
           The uncle of pt is also red
           Only Recoloring required */
        if (uncle && uncle.color === 0) {
          grandParent.color = 0;
          parent.color = 1;
          uncle.color = 1;
          pt = grandParent;
        } else {
          /* Case : 2
             pt is right child of its parent
             Left-rotation required */
          if (pt === parent.right) {
            this.rotateLeft(parent);
            pt = parent;
            parent = pt.parent;
          }

          /* Case : 3
             pt is left child of its parent
             Right-rotation required */
          this.rotateRight(grandParent);
          this.swapColor(parent!, grandParent);
          pt = parent!;
        }
      } else {
        /* Case : B
         Parent of pt is right child of Grand-parent of pt */
        const uncle = grandParent!.left;

        /*  Case : 1
          The uncle of pt is also red
          Only Recoloring required */
        if (uncle != null && uncle.color === 0) {
          grandParent!.color = 0;
          parent.color = 1;
          uncle.color = 1;
          pt = grandParent!;
        } else {
          /* Case : 2
             pt is left child of its parent
             Right-rotation required */
          if (pt === parent.left) {
            this.rotateRight(parent);
            pt = parent;
            parent = pt.parent;
          }

          /* Case : 3
             pt is right child of its parent
             Left-rotation required */
          this.rotateLeft(grandParent!);
          this.swapColor(parent!, grandParent!);
          pt = parent!;
        }
      }
    }
    this.root!.color = 1;
  }

  delete(val: T): boolean {
    const node = this.find(val);
    if (!node) return false;
    node.count--;
    if (!node.count) this.deleteNode(node);
    return true;
  }

  deleteAll(val: T): boolean {
    const node = this.find(val);
    if (!node) return false;
    this.deleteNode(node);
    return true;
  }

  deleteNode(v: RBTreeNode<T>): void {
    const u = BSTreplace(v);

    // True when u and v are both black
    const uvBlack = (u === null || u.color === 1) && v.color === 1;
    const parent = v.parent!;

    if (!u) {
      // u is null therefore v is leaf
      if (v === this.root) this.root = null;
      // v is root, making root null
      else {
        if (uvBlack) {
          // u and v both black
          // v is leaf, fix double black at v
          this.fixDoubleBlack(v);
        } else {
          // u or v is red
          if (v.sibling()) {
            // sibling is not null, make it red"
            v.sibling()!.color = 0;
          }
        }
        // delete v from the tree
        if (v.isOnLeft()) parent.left = null;
        else parent.right = null;
      }
      return;
    }

    if (!v.left || !v.right) {
      // v has 1 child
      if (v === this.root) {
        // v is root, assign the value of u to v, and delete u
        v.data = u.data;
        v.left = v.right = null;
      } else {
        // Detach v from tree and move u up
        if (v.isOnLeft()) parent.left = u;
        else parent.right = u;
        u.parent = parent;
        if (uvBlack) this.fixDoubleBlack(u);
        // u and v both black, fix double black at u
        else u.color = 1; // u or v red, color u black
      }
      return;
    }

    // v has 2 children, swap data with successor and recurse
    this.swapData(u, v);
    this.deleteNode(u);

    // find node that replaces a deleted node in BST
    function BSTreplace(x: RBTreeNode<T>): RBTreeNode<T> | null {
      // when node have 2 children
      if (x.left && x.right) return successor(x.right);
      // when leaf
      if (!x.left && !x.right) return null;
      // when single child
      return x.left ?? x.right;
    }
    // find node that do not have a left child
    // in the subtree of the given node
    function successor(x: RBTreeNode<T>): RBTreeNode<T> {
      let temp = x;
      while (temp.left) temp = temp.left;
      return temp;
    }
  }

  fixDoubleBlack(x: RBTreeNode<T>): void {
    if (x === this.root) return; // Reached root

    const sibling = x.sibling();
    const parent = x.parent!;
    if (!sibling) {
      // No sibiling, double black pushed up
      this.fixDoubleBlack(parent);
    } else {
      if (sibling.color === 0) {
        // Sibling red
        parent.color = 0;
        sibling.color = 1;
        if (sibling.isOnLeft()) this.rotateRight(parent);
        // left case
        else this.rotateLeft(parent); // right case
        this.fixDoubleBlack(x);
      } else {
        // Sibling black
        if (sibling.hasRedChild()) {
          // at least 1 red children
          if (sibling.left && sibling.left.color === 0) {
            if (sibling.isOnLeft()) {
              // left left
              sibling.left.color = sibling.color;
              sibling.color = parent.color;
              this.rotateRight(parent);
            } else {
              // right left
              sibling.left.color = parent.color;
              this.rotateRight(sibling);
              this.rotateLeft(parent);
            }
          } else {
            if (sibling.isOnLeft()) {
              // left right
              sibling.right!.color = parent.color;
              this.rotateLeft(sibling);
              this.rotateRight(parent);
            } else {
              // right right
              sibling.right!.color = sibling.color;
              sibling.color = parent.color;
              this.rotateLeft(parent);
            }
          }
          parent.color = 1;
        } else {
          // 2 black children
          sibling.color = 0;
          if (parent.color === 1) this.fixDoubleBlack(parent);
          else parent.color = 1;
        }
      }
    }
  }

  insert(data: T): boolean {
    // search for a position to insert
    let parent = this.root;
    while (parent) {
      if (this.lt(data, parent.data)) {
        if (!parent.left) break;
        else parent = parent.left;
      } else if (this.lt(parent.data, data)) {
        if (!parent.right) break;
        else parent = parent.right;
      } else break;
    }

    // insert node into parent
    const node = new RBTreeNode(data);
    if (!parent) this.root = node;
    else if (this.lt(node.data, parent.data)) parent.left = node;
    else if (this.lt(parent.data, node.data)) parent.right = node;
    else {
      parent.count++;
      return false;
    }
    node.parent = parent;
    this.fixAfterInsert(node);
    return true;
  }

  find(data: T): RBTreeNode<T> | null {
    let p = this.root;
    while (p) {
      if (this.lt(data, p.data)) {
        p = p.left;
      } else if (this.lt(p.data, data)) {
        p = p.right;
      } else break;
    }
    return p ?? null;
  }

  *inOrder(root: RBTreeNode<T> = this.root!): Generator<T, undefined, void> {
    if (!root) return;
    for (const v of this.inOrder(root.left!)) yield v;
    yield root.data;
    for (const v of this.inOrder(root.right!)) yield v;
  }

  *reverseInOrder(root: RBTreeNode<T> = this.root!): Generator<T, undefined, void> {
    if (!root) return;
    for (const v of this.reverseInOrder(root.right!)) yield v;
    yield root.data;
    for (const v of this.reverseInOrder(root.left!)) yield v;
  }
}

class TreeSet<T = number> {
  _size: number;
  tree: RBTree<T>;
  compare: Compare<T>;
  constructor(
    collection: T[] | Compare<T> = [],
    compare: Compare<T> = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0),
  ) {
    if (typeof collection === 'function') {
      compare = collection;
      collection = [];
    }
    this._size = 0;
    this.compare = compare;
    this.tree = new RBTree(compare);
    for (const val of collection) this.add(val);
  }

  size(): number {
    return this._size;
  }

  has(val: T): boolean {
    return !!this.tree.find(val);
  }

  add(val: T): boolean {
    const successful = this.tree.insert(val);
    this._size += successful ? 1 : 0;
    return successful;
  }

  delete(val: T): boolean {
    const deleted = this.tree.deleteAll(val);
    this._size -= deleted ? 1 : 0;
    return deleted;
  }

  ceil(val: T): T | undefined {
    let p = this.tree.root;
    let higher = null;
    while (p) {
      if (this.compare(p.data, val) >= 0) {
        higher = p;
        p = p.left;
      } else {
        p = p.right;
      }
    }
    return higher?.data;
  }

  floor(val: T): T | undefined {
    let p = this.tree.root;
    let lower = null;
    while (p) {
      if (this.compare(val, p.data) >= 0) {
        lower = p;
        p = p.right;
      } else {
        p = p.left;
      }
    }
    return lower?.data;
  }

  higher(val: T): T | undefined {
    let p = this.tree.root;
    let higher = null;
    while (p) {
      if (this.compare(val, p.data) < 0) {
        higher = p;
        p = p.left;
      } else {
        p = p.right;
      }
    }
    return higher?.data;
  }

  lower(val: T): T | undefined {
    let p = this.tree.root;
    let lower = null;
    while (p) {
      if (this.compare(p.data, val) < 0) {
        lower = p;
        p = p.right;
      } else {
        p = p.left;
      }
    }
    return lower?.data;
  }

  first(): T | undefined {
    return this.tree.inOrder().next().value;
  }

  last(): T | undefined {
    return this.tree.reverseInOrder().next().value;
  }

  shift(): T | undefined {
    const first = this.first();
    if (first === undefined) return undefined;
    this.delete(first);
    return first;
  }

  pop(): T | undefined {
    const last = this.last();
    if (last === undefined) return undefined;
    this.delete(last);
    return last;
  }

  *[Symbol.iterator](): Generator<T, void, void> {
    for (const val of this.values()) yield val;
  }

  *keys(): Generator<T, void, void> {
    for (const val of this.values()) yield val;
  }

  *values(): Generator<T, undefined, void> {
    for (const val of this.tree.inOrder()) yield val;
    return undefined;
  }

  /**
   * Return a generator for reverse order traversing the set
   */
  *rvalues(): Generator<T, undefined, void> {
    for (const val of this.tree.reverseInOrder()) yield val;
    return undefined;
  }
}

class TreeMultiSet<T = number> {
  _size: number;
  tree: RBTree<T>;
  compare: Compare<T>;
  constructor(
    collection: T[] | Compare<T> = [],
    compare: Compare<T> = (l: T, r: T) => (l < r ? -1 : l > r ? 1 : 0),
  ) {
    if (typeof collection === 'function') {
      compare = collection;
      collection = [];
    }
    this._size = 0;
    this.compare = compare;
    this.tree = new RBTree(compare);
    for (const val of collection) this.add(val);
  }

  size(): number {
    return this._size;
  }

  has(val: T): boolean {
    return !!this.tree.find(val);
  }

  add(val: T): boolean {
    const successful = this.tree.insert(val);
    this._size++;
    return successful;
  }

  delete(val: T): boolean {
    const successful = this.tree.delete(val);
    if (!successful) return false;
    this._size--;
    return true;
  }

  count(val: T): number {
    const node = this.tree.find(val);
    return node ? node.count : 0;
  }

  ceil(val: T): T | undefined {
    let p = this.tree.root;
    let higher = null;
    while (p) {
      if (this.compare(p.data, val) >= 0) {
        higher = p;
        p = p.left;
      } else {
        p = p.right;
      }
    }
    return higher?.data;
  }

  floor(val: T): T | undefined {
    let p = this.tree.root;
    let lower = null;
    while (p) {
      if (this.compare(val, p.data) >= 0) {
        lower = p;
        p = p.right;
      } else {
        p = p.left;
      }
    }
    return lower?.data;
  }

  higher(val: T): T | undefined {
    let p = this.tree.root;
    let higher = null;
    while (p) {
      if (this.compare(val, p.data) < 0) {
        higher = p;
        p = p.left;
      } else {
        p = p.right;
      }
    }
    return higher?.data;
  }

  lower(val: T): T | undefined {
    let p = this.tree.root;
    let lower = null;
    while (p) {
      if (this.compare(p.data, val) < 0) {
        lower = p;
        p = p.right;
      } else {
        p = p.left;
      }
    }
    return lower?.data;
  }

  first(): T | undefined {
    return this.tree.inOrder().next().value;
  }

  last(): T | undefined {
    return this.tree.reverseInOrder().next().value;
  }

  shift(): T | undefined {
    const first = this.first();
    if (first === undefined) return undefined;
    this.delete(first);
    return first;
  }

  pop(): T | undefined {
    const last = this.last();
    if (last === undefined) return undefined;
    this.delete(last);
    return last;
  }

  *[Symbol.iterator](): Generator<T, void, void> {
    yield* this.values();
  }

  *keys(): Generator<T, void, void> {
    for (const val of this.values()) yield val;
  }

  *values(): Generator<T, undefined, void> {
    for (const val of this.tree.inOrder()) {
      let count = this.count(val);
      while (count--) yield val;
    }
    return undefined;
  }

  /**
   * Return a generator for reverse order traversing the multi-set
   */
  *rvalues(): Generator<T, undefined, void> {
    for (const val of this.tree.reverseInOrder()) {
      let count = this.count(val);
      while (count--) yield val;
    }
    return undefined;
  }
}

Solution 2

function minReverseOperations(n: number, p: number, banned: number[], k: number): number[] {
  const ans = new Array(n).fill(-1);
  const ts = new Array(2).fill(0).map(() => new TreapMultiSet<number>());
  for (let i = 0; i < n; ++i) {
    ts[i % 2].add(i);
  }
  ans[p] = 0;
  ts[p % 2].delete(p);
  for (const i of banned) {
    ts[i % 2].delete(i);
  }
  ts[0].add(n);
  ts[1].add(n);
  let q = [p];
  while (q.length) {
    const t: number[] = [];
    for (const i of q) {
      const mi = Math.max(i - k + 1, k - i - 1);
      const mx = Math.min(i + k - 1, n * 2 - k - i - 1);
      const s = ts[mi % 2];
      for (let j = s.ceil(mi)!; j <= mx; j = s.ceil(j)!) {
        t.push(j);
        ans[j] = ans[i] + 1;
        s.delete(j);
      }
    }
    q = t;
  }
  return ans;
}

type CompareFunction<T, R extends 'number' | 'boolean'> = (
  a: T,
  b: T,
) => R extends 'number' ? number : boolean;

interface ITreapMultiSet<T> extends Iterable<T> {
  add: (...value: T[]) => this;
  has: (value: T) => boolean;
  delete: (value: T) => void;

  bisectLeft: (value: T) => number;
  bisectRight: (value: T) => number;

  indexOf: (value: T) => number;
  lastIndexOf: (value: T) => number;

  at: (index: number) => T | undefined;
  first: () => T | undefined;
  last: () => T | undefined;

  lower: (value: T) => T | undefined;
  higher: (value: T) => T | undefined;
  floor: (value: T) => T | undefined;
  ceil: (value: T) => T | undefined;

  shift: () => T | undefined;
  pop: (index?: number) => T | undefined;

  count: (value: T) => number;

  keys: () => IterableIterator<T>;
  values: () => IterableIterator<T>;
  rvalues: () => IterableIterator<T>;
  entries: () => IterableIterator<[number, T]>;

  readonly size: number;
}

class TreapNode<T = number> {
  value: T;
  count: number;
  size: number;
  priority: number;
  left: TreapNode<T> | null;
  right: TreapNode<T> | null;

  constructor(value: T) {
    this.value = value;
    this.count = 1;
    this.size = 1;
    this.priority = Math.random();
    this.left = null;
    this.right = null;
  }

  static getSize(node: TreapNode<any> | null): number {
    return node?.size ?? 0;
  }

  static getFac(node: TreapNode<any> | null): number {
    return node?.priority ?? 0;
  }

  pushUp(): void {
    let tmp = this.count;
    tmp += TreapNode.getSize(this.left);
    tmp += TreapNode.getSize(this.right);
    this.size = tmp;
  }

  rotateRight(): TreapNode<T> {
    // eslint-disable-next-line @typescript-eslint/no-this-alias
    let node: TreapNode<T> = this;
    const left = node.left;
    node.left = left?.right ?? null;
    left && (left.right = node);
    left && (node = left);
    node.right?.pushUp();
    node.pushUp();
    return node;
  }

  rotateLeft(): TreapNode<T> {
    // eslint-disable-next-line @typescript-eslint/no-this-alias
    let node: TreapNode<T> = this;
    const right = node.right;
    node.right = right?.left ?? null;
    right && (right.left = node);
    right && (node = right);
    node.left?.pushUp();
    node.pushUp();
    return node;
  }
}

class TreapMultiSet<T = number> implements ITreapMultiSet<T> {
  private readonly root: TreapNode<T>;
  private readonly compareFn: CompareFunction<T, 'number'>;
  private readonly leftBound: T;
  private readonly rightBound: T;

  constructor(compareFn?: CompareFunction<T, 'number'>);
  constructor(compareFn: CompareFunction<T, 'number'>, leftBound: T, rightBound: T);
  constructor(
    compareFn: CompareFunction<T, any> = (a: any, b: any) => a - b,
    leftBound: any = -Infinity,
    rightBound: any = Infinity,
  ) {
    this.root = new TreapNode<T>(rightBound);
    this.root.priority = Infinity;
    this.root.left = new TreapNode<T>(leftBound);
    this.root.left.priority = -Infinity;
    this.root.pushUp();

    this.leftBound = leftBound;
    this.rightBound = rightBound;
    this.compareFn = compareFn;
  }

  get size(): number {
    return this.root.size - 2;
  }

  get height(): number {
    const getHeight = (node: TreapNode<T> | null): number => {
      if (node == null) return 0;
      return 1 + Math.max(getHeight(node.left), getHeight(node.right));
    };

    return getHeight(this.root);
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Returns true if value is a member.
   */
  has(value: T): boolean {
    const compare = this.compareFn;
    const dfs = (node: TreapNode<T> | null, value: T): boolean => {
      if (node == null) return false;
      if (compare(node.value, value) === 0) return true;
      if (compare(node.value, value) < 0) return dfs(node.right, value);
      return dfs(node.left, value);
    };

    return dfs(this.root, value);
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Add value to sorted set.
   */
  add(...values: T[]): this {
    const compare = this.compareFn;
    const dfs = (
      node: TreapNode<T> | null,
      value: T,
      parent: TreapNode<T>,
      direction: 'left' | 'right',
    ): void => {
      if (node == null) return;
      if (compare(node.value, value) === 0) {
        node.count++;
        node.pushUp();
      } else if (compare(node.value, value) > 0) {
        if (node.left) {
          dfs(node.left, value, node, 'left');
        } else {
          node.left = new TreapNode(value);
          node.pushUp();
        }

        if (TreapNode.getFac(node.left) > node.priority) {
          parent[direction] = node.rotateRight();
        }
      } else if (compare(node.value, value) < 0) {
        if (node.right) {
          dfs(node.right, value, node, 'right');
        } else {
          node.right = new TreapNode(value);
          node.pushUp();
        }

        if (TreapNode.getFac(node.right) > node.priority) {
          parent[direction] = node.rotateLeft();
        }
      }
      parent.pushUp();
    };

    values.forEach(value => dfs(this.root.left, value, this.root, 'left'));
    return this;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Remove value from sorted set if it is a member.
   * If value is not a member, do nothing.
   */
  delete(value: T): void {
    const compare = this.compareFn;
    const dfs = (
      node: TreapNode<T> | null,
      value: T,
      parent: TreapNode<T>,
      direction: 'left' | 'right',
    ): void => {
      if (node == null) return;

      if (compare(node.value, value) === 0) {
        if (node.count > 1) {
          node.count--;
          node?.pushUp();
        } else if (node.left == null && node.right == null) {
          parent[direction] = null;
        } else {
          // 旋到根节点
          if (
            node.right == null ||
            TreapNode.getFac(node.left) > TreapNode.getFac(node.right)
          ) {
            parent[direction] = node.rotateRight();
            dfs(parent[direction]?.right ?? null, value, parent[direction]!, 'right');
          } else {
            parent[direction] = node.rotateLeft();
            dfs(parent[direction]?.left ?? null, value, parent[direction]!, 'left');
          }
        }
      } else if (compare(node.value, value) > 0) {
        dfs(node.left, value, node, 'left');
      } else if (compare(node.value, value) < 0) {
        dfs(node.right, value, node, 'right');
      }

      parent?.pushUp();
    };

    dfs(this.root.left, value, this.root, 'left');
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Returns an index to insert value in the sorted set.
   * If the value is already present, the insertion point will be before (to the left of) any existing values.
   */
  bisectLeft(value: T): number {
    const compare = this.compareFn;
    const dfs = (node: TreapNode<T> | null, value: T): number => {
      if (node == null) return 0;

      if (compare(node.value, value) === 0) {
        return TreapNode.getSize(node.left);
      } else if (compare(node.value, value) > 0) {
        return dfs(node.left, value);
      } else if (compare(node.value, value) < 0) {
        return dfs(node.right, value) + TreapNode.getSize(node.left) + node.count;
      }

      return 0;
    };

    return dfs(this.root, value) - 1;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Returns an index to insert value in the sorted set.
   * If the value is already present, the insertion point will be before (to the right of) any existing values.
   */
  bisectRight(value: T): number {
    const compare = this.compareFn;
    const dfs = (node: TreapNode<T> | null, value: T): number => {
      if (node == null) return 0;

      if (compare(node.value, value) === 0) {
        return TreapNode.getSize(node.left) + node.count;
      } else if (compare(node.value, value) > 0) {
        return dfs(node.left, value);
      } else if (compare(node.value, value) < 0) {
        return dfs(node.right, value) + TreapNode.getSize(node.left) + node.count;
      }

      return 0;
    };
    return dfs(this.root, value) - 1;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Returns the index of the first occurrence of a value in the set, or -1 if it is not present.
   */
  indexOf(value: T): number {
    const compare = this.compareFn;
    let isExist = false;

    const dfs = (node: TreapNode<T> | null, value: T): number => {
      if (node == null) return 0;

      if (compare(node.value, value) === 0) {
        isExist = true;
        return TreapNode.getSize(node.left);
      } else if (compare(node.value, value) > 0) {
        return dfs(node.left, value);
      } else if (compare(node.value, value) < 0) {
        return dfs(node.right, value) + TreapNode.getSize(node.left) + node.count;
      }

      return 0;
    };
    const res = dfs(this.root, value) - 1;
    return isExist ? res : -1;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Returns the index of the last occurrence of a value in the set, or -1 if it is not present.
   */
  lastIndexOf(value: T): number {
    const compare = this.compareFn;
    let isExist = false;

    const dfs = (node: TreapNode<T> | null, value: T): number => {
      if (node == null) return 0;

      if (compare(node.value, value) === 0) {
        isExist = true;
        return TreapNode.getSize(node.left) + node.count - 1;
      } else if (compare(node.value, value) > 0) {
        return dfs(node.left, value);
      } else if (compare(node.value, value) < 0) {
        return dfs(node.right, value) + TreapNode.getSize(node.left) + node.count;
      }

      return 0;
    };

    const res = dfs(this.root, value) - 1;
    return isExist ? res : -1;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Returns the item located at the specified index.
   * @param index The zero-based index of the desired code unit. A negative index will count back from the last item.
   */
  at(index: number): T | undefined {
    if (index < 0) index += this.size;
    if (index < 0 || index >= this.size) return undefined;

    const dfs = (node: TreapNode<T> | null, rank: number): T | undefined => {
      if (node == null) return undefined;

      if (TreapNode.getSize(node.left) >= rank) {
        return dfs(node.left, rank);
      } else if (TreapNode.getSize(node.left) + node.count >= rank) {
        return node.value;
      } else {
        return dfs(node.right, rank - TreapNode.getSize(node.left) - node.count);
      }
    };

    const res = dfs(this.root, index + 2);
    return ([this.leftBound, this.rightBound] as any[]).includes(res) ? undefined : res;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Find and return the element less than `val`, return `undefined` if no such element found.
   */
  lower(value: T): T | undefined {
    const compare = this.compareFn;
    const dfs = (node: TreapNode<T> | null, value: T): T | undefined => {
      if (node == null) return undefined;
      if (compare(node.value, value) >= 0) return dfs(node.left, value);

      const tmp = dfs(node.right, value);
      if (tmp == null || compare(node.value, tmp) > 0) {
        return node.value;
      } else {
        return tmp;
      }
    };

    const res = dfs(this.root, value) as any;
    return res === this.leftBound ? undefined : res;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Find and return the element greater than `val`, return `undefined` if no such element found.
   */
  higher(value: T): T | undefined {
    const compare = this.compareFn;
    const dfs = (node: TreapNode<T> | null, value: T): T | undefined => {
      if (node == null) return undefined;
      if (compare(node.value, value) <= 0) return dfs(node.right, value);

      const tmp = dfs(node.left, value);

      if (tmp == null || compare(node.value, tmp) < 0) {
        return node.value;
      } else {
        return tmp;
      }
    };

    const res = dfs(this.root, value) as any;
    return res === this.rightBound ? undefined : res;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Find and return the element less than or equal to `val`, return `undefined` if no such element found.
   */
  floor(value: T): T | undefined {
    const compare = this.compareFn;
    const dfs = (node: TreapNode<T> | null, value: T): T | undefined => {
      if (node == null) return undefined;
      if (compare(node.value, value) === 0) return node.value;
      if (compare(node.value, value) >= 0) return dfs(node.left, value);

      const tmp = dfs(node.right, value);
      if (tmp == null || compare(node.value, tmp) > 0) {
        return node.value;
      } else {
        return tmp;
      }
    };

    const res = dfs(this.root, value) as any;
    return res === this.leftBound ? undefined : res;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description Find and return the element greater than or equal to `val`, return `undefined` if no such element found.
   */
  ceil(value: T): T | undefined {
    const compare = this.compareFn;
    const dfs = (node: TreapNode<T> | null, value: T): T | undefined => {
      if (node == null) return undefined;
      if (compare(node.value, value) === 0) return node.value;
      if (compare(node.value, value) <= 0) return dfs(node.right, value);

      const tmp = dfs(node.left, value);

      if (tmp == null || compare(node.value, tmp) < 0) {
        return node.value;
      } else {
        return tmp;
      }
    };

    const res = dfs(this.root, value) as any;
    return res === this.rightBound ? undefined : res;
  }

  /**
   * @complexity `O(logn)`
   * @description
   * Returns the last element from set.
   * If the set is empty, undefined is returned.
   */
  first(): T | undefined {
    const iter = this.inOrder();
    iter.next();
    const res = iter.next().value;
    return res === this.rightBound ? undefined : res;
  }

  /**
   * @complexity `O(logn)`
   * @description
   * Returns the last element from set.
   * If the set is empty, undefined is returned .
   */
  last(): T | undefined {
    const iter = this.reverseInOrder();
    iter.next();
    const res = iter.next().value;
    return res === this.leftBound ? undefined : res;
  }

  /**
   * @complexity `O(logn)`
   * @description
   * Removes the first element from an set and returns it.
   * If the set is empty, undefined is returned and the set is not modified.
   */
  shift(): T | undefined {
    const first = this.first();
    if (first === undefined) return undefined;
    this.delete(first);
    return first;
  }

  /**
   * @complexity `O(logn)`
   * @description
   * Removes the last element from an set and returns it.
   * If the set is empty, undefined is returned and the set is not modified.
   */
  pop(index?: number): T | undefined {
    if (index == null) {
      const last = this.last();
      if (last === undefined) return undefined;
      this.delete(last);
      return last;
    }

    const toDelete = this.at(index);
    if (toDelete == null) return;
    this.delete(toDelete);
    return toDelete;
  }

  /**
   *
   * @complexity `O(logn)`
   * @description
   * Returns number of occurrences of value in the sorted set.
   */
  count(value: T): number {
    const compare = this.compareFn;
    const dfs = (node: TreapNode<T> | null, value: T): number => {
      if (node == null) return 0;
      if (compare(node.value, value) === 0) return node.count;
      if (compare(node.value, value) < 0) return dfs(node.right, value);
      return dfs(node.left, value);
    };

    return dfs(this.root, value);
  }

  *[Symbol.iterator](): Generator<T, any, any> {
    yield* this.values();
  }

  /**
   * @description
   * Returns an iterable of keys in the set.
   */
  *keys(): Generator<T, any, any> {
    yield* this.values();
  }

  /**
   * @description
   * Returns an iterable of values in the set.
   */
  *values(): Generator<T, any, any> {
    const iter = this.inOrder();
    iter.next();
    const steps = this.size;
    for (let _ = 0; _ < steps; _++) {
      yield iter.next().value;
    }
  }

  /**
   * @description
   * Returns a generator for reversed order traversing the set.
   */
  *rvalues(): Generator<T, any, any> {
    const iter = this.reverseInOrder();
    iter.next();
    const steps = this.size;
    for (let _ = 0; _ < steps; _++) {
      yield iter.next().value;
    }
  }

  /**
   * @description
   * Returns an iterable of key, value pairs for every entry in the set.
   */
  *entries(): IterableIterator<[number, T]> {
    const iter = this.inOrder();
    iter.next();
    const steps = this.size;
    for (let i = 0; i < steps; i++) {
      yield [i, iter.next().value];
    }
  }

  private *inOrder(root: TreapNode<T> | null = this.root): Generator<T, any, any> {
    if (root == null) return;
    yield* this.inOrder(root.left);
    const count = root.count;
    for (let _ = 0; _ < count; _++) {
      yield root.value;
    }
    yield* this.inOrder(root.right);
  }

  private *reverseInOrder(root: TreapNode<T> | null = this.root): Generator<T, any, any> {
    if (root == null) return;
    yield* this.reverseInOrder(root.right);
    const count = root.count;
    for (let _ = 0; _ < count; _++) {
      yield root.value;
    }
    yield* this.reverseInOrder(root.left);
  }
}

如果你对这篇内容有疑问,欢迎到本站社区发帖提问 参与讨论,获取更多帮助,或者扫码二维码加入 Web 技术交流群。

扫码二维码加入Web技术交流群

发布评论

需要 登录 才能够评论, 你可以免费 注册 一个本站的账号。
列表为空,暂无数据
    我们使用 Cookies 和其他技术来定制您的体验包括您的登录状态等。通过阅读我们的 隐私政策 了解更多相关信息。 单击 接受 或继续使用网站,即表示您同意使用 Cookies 和您的相关数据。
    原文