Question

2555. 两个线段获得的最多奖品

English Version

题目描述

在 X轴 上有一些奖品。给你一个整数数组 prizePositions ，它按照 非递减 顺序排列，其中 prizePositions[i] 是第 i 件奖品的位置。数轴上一个位置可能会有多件奖品。再给你一个整数 k 。

你可以选择两个端点为整数的线段。每个线段的长度都必须是 k 。你可以获得位置在任一线段上的所有奖品（包括线段的两个端点）。注意，两个线段可能会有相交。

• 比方说 k = 2 ，你可以选择线段 [1, 3] 和 [2, 4] ，你可以获得满足 1 <= prizePositions[i] <= 3 或者 2 <= prizePositions[i] <= 4 的所有奖品 i 。

请你返回在选择两个最优线段的前提下，可以获得的 最多 奖品数目。

 

示例 1：

输入：prizePositions = [1,1,2,2,3,3,5], k = 2
输出：7
解释：这个例子中，你可以选择线段 [1, 3] 和 [3, 5] ，获得 7 个奖品。

示例 2：

输入：prizePositions = [1,2,3,4], k = 0
输出：2
解释：这个例子中，一个选择是选择线段 [3, 3] 和 [4, 4] ，获得 2 个奖品。

 

提示：

• 1 <= prizePositions.length <= 105
• 1 <= prizePositions[i] <= 109
• 0 <= k <= 109
• prizePositions 有序非递减。

解法

方法一：动态规划 + 二分查找

我们定义 $f[i]$ 表示在前 $i$ 个奖品中，选择一个长度为 $k$ 的线段，可以获得的最多奖品数目。初始时 $f[0] = 0$。定义答案变量 $ans = 0$。

接下来，我们枚举每个奖品的位置 $x$，通过二分查找，找到最左边的奖品下标 $j$，使得 $prizePositions[j] \geq x - k$。此时，我们更新答案 $ans = \max(ans, f[j] + i - j)$，并更新 $f[i] = \max(f[i - 1], i - j)$。

最后，返回 $ans$ 即可。

时间复杂度 $O(n \times \log n)$，空间复杂度 $O(n)$。其中 $n$ 为奖品数目。

class Solution:
  def maximizeWin(self, prizePositions: List[int], k: int) -> int:
    n = len(prizePositions)
    f = [0] * (n + 1)
    ans = 0
    for i, x in enumerate(prizePositions, 1):
      j = bisect_left(prizePositions, x - k)
      ans = max(ans, f[j] + i - j)
      f[i] = max(f[i - 1], i - j)
    return ans

class Solution {
  public int maximizeWin(int[] prizePositions, int k) {
    int n = prizePositions.length;
    int[] f = new int[n + 1];
    int ans = 0;
    for (int i = 1; i <= n; ++i) {
      int x = prizePositions[i - 1];
      int j = search(prizePositions, x - k);
      ans = Math.max(ans, f[j] + i - j);
      f[i] = Math.max(f[i - 1], i - j);
    }
    return ans;
  }

  private int search(int[] nums, int x) {
    int left = 0, right = nums.length;
    while (left < right) {
      int mid = (left + right) >> 1;
      if (nums[mid] >= x) {
        right = mid;
      } else {
        left = mid + 1;
      }
    }
    return left;
  }
}

class Solution {
public:
  int maximizeWin(vector<int>& prizePositions, int k) {
    int n = prizePositions.size();
    vector<int> f(n + 1);
    int ans = 0;
    for (int i = 1; i <= n; ++i) {
      int x = prizePositions[i - 1];
      int j = lower_bound(prizePositions.begin(), prizePositions.end(), x - k) - prizePositions.begin();
      ans = max(ans, f[j] + i - j);
      f[i] = max(f[i - 1], i - j);
    }
    return ans;
  }
};

func maximizeWin(prizePositions []int, k int) (ans int) {
  n := len(prizePositions)
  f := make([]int, n+1)
  for i, x := range prizePositions {
    j := sort.Search(n, func(h int) bool { return prizePositions[h] >= x-k })
    ans = max(ans, f[j]+i-j+1)
    f[i+1] = max(f[i], i-j+1)
  }
  return
}

function maximizeWin(prizePositions: number[], k: number): number {
  const n = prizePositions.length;
  const f: number[] = Array(n + 1).fill(0);
  let ans = 0;
  const search = (x: number): number => {
    let left = 0;
    let right = n;
    while (left < right) {
      const mid = (left + right) >> 1;
      if (prizePositions[mid] >= x) {
        right = mid;
      } else {
        left = mid + 1;
      }
    }
    return left;
  };
  for (let i = 1; i <= n; ++i) {
    const x = prizePositions[i - 1];
    const j = search(x - k);
    ans = Math.max(ans, f[j] + i - j);
    f[i] = Math.max(f[i - 1], i - j);
  }
  return ans;
}