Skip to content

2122. Recover the Original Array 👍

  • Time: $O(n^2)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class Solution {
 public:
  vector<int> recoverArray(vector<int>& nums) {
    const int n = nums.size();
    unordered_map<int, int> count;

    for (const int num : nums)
      ++count[num];

    sort(begin(nums), end(nums));

    for (int i = 1; i < n; ++i) {
      const int x = nums[i] - nums[0];  // 2 * k
      if (x <= 0 || x & 1)
        continue;
      vector<int> A = getArray(nums, x, count);
      if (!A.empty())
        return A;
    }

    throw;
  }

 private:
  vector<int> getArray(const vector<int>& nums, int x,
                       unordered_map<int, int> count) {
    vector<int> A;
    for (const int num : nums) {
      if (!count.count(num) || count[num] == 0)
        continue;
      if (!count.count(num + x) || count[num + x] == 0)
        return {};
      --count[num];
      --count[num + x];
      A.push_back(num + x / 2);
    }
    return A;
  }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class Solution {
  public int[] recoverArray(int[] nums) {
    final int n = nums.length;
    Map<Integer, Integer> count = new HashMap<>();

    for (final int num : nums)
      count.merge(num, 1, Integer::sum);

    Arrays.sort(nums);

    for (int i = 1; i < n; ++i) {
      final int x = nums[i] - nums[0]; // 2 * k
      if (x <= 0 || x % 2 == 1)
        continue;
      Map<Integer, Integer> countCopy = new HashMap<>();
      countCopy.putAll(count);
      int[] A = getArray(nums, x, countCopy);
      if (A.length == n / 2)
        return A;
    }

    throw new IllegalArgumentException();
  }

  private int[] getArray(int[] nums, int x, Map<Integer, Integer> count) {
    List<Integer> A = new ArrayList<>();
    for (final int num : nums) {
      if (count.getOrDefault(num, 0) == 0)
        continue;
      if (count.getOrDefault(num + x, 0) == 0)
        return new int[] {};
      count.merge(num, -1, Integer::sum);
      count.merge(num + x, -1, Integer::sum);
      A.add(num + x / 2);
    }
    return A.stream().mapToInt(i -> i).toArray();
  }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution:
  def recoverArray(self, nums: List[int]) -> List[int]:
    nums = sorted(nums)

    def getArray(x: int, count: collections.Counter) -> List[int]:
      A = []
      for num in nums:
        if count[num] == 0:
          continue
        if count[num + x] == 0:
          return []
        count[num] -= 1
        count[num + x] -= 1
        A.append(num + x // 2)
      return A

    count = Counter(nums)

    for i in range(1, len(nums)):
      x = nums[i] - nums[0]  # 2 * k
      if x <= 0 or x & 1:
        continue
      A = getArray(x, count.copy())
      if A:
        return A