Skip to content

2731. Movement of Robots 👍

  • Time: $O(\texttt{sort})$
  • 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
class Solution {
 public:
  int sumDistance(vector<int>& nums, string s, int d) {
    constexpr int kMod = 1'000'000'007;
    const int n = nums.size();
    int ans = 0;
    int prefix = 0;
    vector<int> pos;

    for (int i = 0; i < nums.size(); ++i)
      if (s[i] == 'L')
        pos.push_back(nums[i] - d);
      else
        pos.push_back(nums[i] + d);

    ranges::sort(pos);

    for (int i = 0; i < n; ++i) {
      ans =
          ((ans + static_cast<long>(i) * pos[i] - prefix) % kMod + kMod) % kMod;
      prefix = ((0L + prefix + pos[i]) % kMod + kMod) % kMod;
    }

    return ans;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
  public int sumDistance(int[] nums, String s, int d) {
    final int kMod = 1_000_000_007;
    final int n = nums.length;
    int ans = 0;
    int prefix = 0;
    int[] pos = new int[n];

    for (int i = 0; i < n; ++i)
      if (s.charAt(i) == 'L')
        pos[i] = nums[i] - d;
      else
        pos[i] = nums[i] + d;

    Arrays.sort(pos);

    for (int i = 0; i < n; ++i) {
      ans = (int) (((ans + 1L * i * pos[i] - prefix) % kMod + kMod) % kMod);
      prefix = (int) (((0L + prefix + pos[i]) % kMod + kMod) % kMod);
    }

    return ans;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution:
  def sumDistance(self, nums: list[int], s: str, d: int) -> int:
    kMod = 1_000_000_007
    ans = 0
    prefix = 0
    pos = sorted([num - d if c == 'L' else num + d
                  for num, c in zip(nums, s)])

    for i, p in enumerate(pos):
      ans = ((ans + i * p - prefix) % kMod + kMod) % kMod
      prefix = ((prefix + p) % kMod + kMod) % kMod

    return ans