3041. Maximize Consecutive Elements in an Array After Modification

Approach 1: $O(n)$ space

• Time: $O(\texttt{sort})$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  int maxSelectedElements(vector<int>& nums) {
    int ans = 0;
    // {num: the length of the longest consecutive elements ending in num}
    unordered_map<int, int> dp;

    ranges::sort(nums);

    for (const int num : nums) {
      dp[num + 1] = dp[num] + 1;
      dp[num] = dp[num - 1] + 1;
      ans = max({ans, dp[num], dp[num + 1]});
    }

    return ans;
  }
};

class Solution {
  public int maxSelectedElements(int[] nums) {
    int ans = 0;
    // {num: the length of the longest consecutive elements ending in num}
    HashMap<Integer, Integer> dp = new HashMap<>();

    Arrays.sort(nums);

    for (final int num : nums) {
      dp.put(num + 1, dp.getOrDefault(num, 0) + 1);
      dp.put(num, dp.getOrDefault(num - 1, 0) + 1);
      ans = Math.max(ans, Math.max(dp.get(num), dp.get(num + 1)));
    }

    return ans;
  }
}

class Solution:
  def maxSelectedElements(self, nums: list[int]) -> int:
    ans = 0
    # {num: the length of the longest consecutive elements ending in num}
    dp = {}

    for num in sorted(nums):
      dp[num + 1] = dp.get(num, 0) + 1
      dp[num] = dp.get(num - 1, 0) + 1
      ans = max(ans, dp[num], dp[num + 1])

    return ans

Approach 2: $O(1)$ space

• Time: $O(\texttt{sort})$
• Space: $O(1)$

C++JavaPython

class Solution {
 public:
  int maxSelectedElements(vector<int>& nums) {
    int ans = 1;
    int prev = INT_MIN;
    // the length of the longest consecutive elements (seq0) ending in the
    // previous number
    int dp0 = 1;
    // the length of the longest consecutive elements (seq1) ending in the
    // previous number + 1
    int dp1 = 1;

    ranges::sort(nums);

    for (const int num : nums) {
      if (num == prev) {
        dp1 = dp0 + 1;  // Append `num + 1` to seq0.
      } else if (num == prev + 1) {
        ++dp0;  // Append `num` to seq0.
        ++dp1;  // Add 1 to every number in seq0 and append `num + 1` to seq0.
      } else if (num == prev + 2) {
        dp0 = dp1 + 1;  // Append `num` to seq1.
        dp1 = 1;        // Start a new sequence [`num + 1`].
      } else {
        dp0 = 1;  // Start a new sequence [`num`].
        dp1 = 1;  // Start a new sequence [`num + 1`].
      }
      ans = max({ans, dp0, dp1});
      prev = num;
    }

    return ans;
  }
};

class Solution {
  public int maxSelectedElements(int[] nums) {
    int ans = 1;
    int prev = Integer.MIN_VALUE;
    // the length of the longest consecutive elements (seq0) ending in the
    // previous number
    int dp0 = 1;
    // the length of the longest consecutive elements (seq1) ending in the
    // previous number + 1
    int dp1 = 1;

    Arrays.sort(nums);

    for (final int num : nums) {
      if (num == prev) {
        dp1 = dp0 + 1; // Append `num + 1` to seq0.
      } else if (num == prev + 1) {
        ++dp0; // Append `num` to seq0.
        ++dp1; // Add 1 to every number in seq0 and append `num + 1` to seq0.
      } else if (num == prev + 2) {
        dp0 = dp1 + 1; // Append `num` to seq1.
        dp1 = 1;       // Start a new sequence [`num + 1`].
      } else {
        dp0 = 1; // Start a new sequence [`num`].
        dp1 = 1; // Start a new sequence [`num + 1`].
      }
      ans = Math.max(ans, Math.max(dp0, dp1));
      prev = num;
    }

    return ans;
  }
}

class Solution:
  def maxSelectedElements(self, nums: list[int]) -> int:
    ans = 1
    prev = -math.inf
    # the length of the longest consecutive elements (seq0) ending in the
    # previous number
    dp0 = 1
    # the length of the longest consecutive elements (seq1) ending in the
    # previous number + 1
    dp1 = 1

    for num in sorted(nums):
      if num == prev:
        dp1 = dp0 + 1  # Append `num + 1` to seq0.
      elif num == prev + 1:
        dp0 += 1  # Append `num` to seq0.
        dp1 += 1  # Add 1 to every number in seq0 and append `num + 1` to seq0.
      elif num == prev + 2:
        dp0 = dp1 + 1  # Append `num` to seq1.
        dp1 = 1        # Start a new sequence [`num + 1`].
      else:
        dp0 = 1  # Start a new sequence [`num`].
        dp1 = 1  # Start a new sequence [`num + 1`].
      ans = max(ans, dp0, dp1)
      prev = num

    return ans