Skip to content

3072. Distribute Elements Into Two Arrays II 👍

  • Time: $O(n\log n)$
  • 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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
class FenwickTree {
 public:
  FenwickTree(int n) : sums(n + 1) {}

  void add(int i, int delta) {
    while (i < sums.size()) {
      sums[i] += delta;
      i += lowbit(i);
    }
  }

  int get(int i) const {
    int sum = 0;
    while (i > 0) {
      sum += sums[i];
      i -= lowbit(i);
    }
    return sum;
  }

 private:
  vector<int> sums;

  static inline int lowbit(int i) {
    return i & -i;
  }
};

class Solution {
 public:
  vector<int> resultArray(vector<int>& nums) {
    vector<int> arr1;
    vector<int> arr2;
    const unordered_map<int, int> ranks = getRanks(nums);
    FenwickTree tree1(ranks.size());
    FenwickTree tree2(ranks.size());

    add(nums[0], arr1, tree1, ranks);
    add(nums[1], arr2, tree2, ranks);

    for (int i = 2; i < nums.size(); ++i) {
      const int greaterCount1 = arr1.size() - tree1.get(ranks.at(nums[i]));
      const int greaterCount2 = arr2.size() - tree2.get(ranks.at(nums[i]));
      if (greaterCount1 > greaterCount2)
        add(nums[i], arr1, tree1, ranks);
      else if (greaterCount1 < greaterCount2)
        add(nums[i], arr2, tree2, ranks);
      else if (arr1.size() > arr2.size())
        add(nums[i], arr2, tree2, ranks);
      else
        add(nums[i], arr1, tree1, ranks);
    }

    arr1.insert(arr1.end(), arr2.begin(), arr2.end());
    return arr1;
  }

 private:
  unordered_map<int, int> getRanks(const vector<int>& nums) {
    unordered_map<int, int> ranks;
    set<int> sorted(nums.begin(), nums.end());
    int rank = 0;
    for (const int num : sorted)
      ranks[num] = ++rank;
    return ranks;
  }

  void add(int num, vector<int>& arr, FenwickTree& tree,
           const unordered_map<int, int>& ranks) {
    arr.push_back(num);
    tree.add(ranks.at(num), 1);
  };
};

 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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
class FenwickTree {
  public FenwickTree(int n) {
    sums = new int[n + 1];
  }

  public void add(int i, int delta) {
    while (i < sums.length) {
      sums[i] += delta;
      i += lowbit(i);
    }
  }

  public int get(int i) {
    int sum = 0;
    while (i > 0) {
      sum += sums[i];
      i -= lowbit(i);
    }
    return sum;
  }

  private int[] sums;

  private static int lowbit(int i) {
    return i & -i;
  }
}

class Solution {
  public int[] resultArray(int[] nums) {
    List<Integer> arr1 = new ArrayList<>();
    List<Integer> arr2 = new ArrayList<>();
    Map<Integer, Integer> ranks = getRanks(nums);
    FenwickTree tree1 = new FenwickTree(ranks.size());
    FenwickTree tree2 = new FenwickTree(ranks.size());

    add(nums[0], arr1, tree1, ranks);
    add(nums[1], arr2, tree2, ranks);

    for (int i = 2; i < nums.length; ++i) {
      final int greaterCount1 = arr1.size() - tree1.get(ranks.get(nums[i]));
      final int greaterCount2 = arr2.size() - tree2.get(ranks.get(nums[i]));
      if (greaterCount1 > greaterCount2)
        add(nums[i], arr1, tree1, ranks);
      else if (greaterCount1 < greaterCount2)
        add(nums[i], arr2, tree2, ranks);
      else if (arr1.size() > arr2.size())
        add(nums[i], arr2, tree2, ranks);
      else
        add(nums[i], arr1, tree1, ranks);
    }

    arr1.addAll(arr2);
    return arr1.stream().mapToInt(Integer::intValue).toArray();
  }

  private Map<Integer, Integer> getRanks(int[] nums) {
    Map<Integer, Integer> ranks = new HashMap<>();
    SortedSet<Integer> sorted = new TreeSet<>();
    for (final int num : nums)
      sorted.add(num);
    int rank = 0;
    for (Iterator<Integer> it = sorted.iterator(); it.hasNext();)
      ranks.put(it.next(), ++rank);
    return ranks;
  }

  private void add(int num, List<Integer> arr, FenwickTree tree, Map<Integer, Integer> ranks) {
    arr.add(num);
    tree.add(ranks.get(num), 1);
  };
}

 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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
class FenwickTree:
  def __init__(self, n: int):
    self.sums = [0] * (n + 1)

  def add(self, i: int, delta: int) -> None:
    while i < len(self.sums):
      self.sums[i] += delta
      i += FenwickTree.lowbit(i)

  def get(self, i: int) -> int:
    summ = 0
    while i > 0:
      summ += self.sums[i]
      i -= FenwickTree.lowbit(i)
    return summ

  @staticmethod
  def lowbit(i: int) -> int:
    return i & -i


class Solution:
  def resultArray(self, nums: list[int]) -> list[int]:
    arr1 = []
    arr2 = []
    ranks = self._getRanks(nums)
    tree1 = FenwickTree(len(ranks))
    tree2 = FenwickTree(len(ranks))

    def add(num: int, arr: list[int], tree: FenwickTree) -> None:
      arr.append(num)
      tree.add(ranks[num], 1)

    add(nums[0], arr1, tree1)
    add(nums[1], arr2, tree2)

    for i in range(2, len(nums)):
      greaterCount1 = len(arr1) - tree1.get(ranks[nums[i]])
      greaterCount2 = len(arr2) - tree2.get(ranks[nums[i]])
      if greaterCount1 > greaterCount2:
        add(nums[i], arr1, tree1)
      elif greaterCount1 < greaterCount2:
        add(nums[i], arr2, tree2)
      elif len(arr1) > len(arr2):
        add(nums[i], arr2, tree2)
      else:
        add(nums[i], arr1, tree1)

    return arr1 + arr2

  def _getRanks(self, nums: list[int]) -> dict[int, int]:
    ranks = collections.Counter()
    rank = 0
    for num in sorted(set(nums)):
      rank += 1
      ranks[num] = rank
    return ranks