2463. Minimum Total Distance Traveled

• Time: $O(|\texttt{robot}|^2|\texttt{factory}|)$
• Space: $O(|\texttt{robot}|^2|\texttt{factory}|)$

C++JavaPython

class Solution {
 public:
  long long minimumTotalDistance(vector<int>& robot,
                                 vector<vector<int>>& factory) {
    ranges::sort(robot);
    ranges::sort(factory);
    vector<vector<vector<long>>> mem(
        robot.size(),
        vector<vector<long>>(factory.size(), vector<long>(robot.size())));
    return minimumTotalDistance(robot, factory, 0, 0, 0, mem);
  }

 private:
  // Returns the minimum distance to fix robot[i..n) with factory[j..n), where
  // factory[j] already fixed k robots.
  long minimumTotalDistance(const vector<int>& robot,
                            const vector<vector<int>>& factory, int i, int j,
                            int k, vector<vector<vector<long>>>& mem) {
    if (i == robot.size())
      return 0;
    if (j == factory.size())
      return LONG_MAX;
    if (mem[i][j][k] > 0)
      return mem[i][j][k];
    const long skipFactory =
        minimumTotalDistance(robot, factory, i, j + 1, 0, mem);
    const int position = factory[j][0];
    const int limit = factory[j][1];
    const long useFactory =
        limit > k ? minimumTotalDistance(robot, factory, i + 1, j, k + 1, mem) +
                        abs(robot[i] - position)
                  : LONG_MAX / 2;
    return mem[i][j][k] = min(skipFactory, useFactory);
  }
};

class Solution {
  public long minimumTotalDistance(List<Integer> robot, int[][] factory) {
    Collections.sort(robot);
    Arrays.sort(factory, (a, b) -> Integer.compare(a[0], b[0]));
    long[][][] mem = new long[robot.size()][factory.length][robot.size()];
    return minimumTotalDistance(robot, factory, 0, 0, 0, mem);
  }

  private long minimumTotalDistance(List<Integer> robot, int[][] factory, int i, int j, int k,
                                    long[][][] mem) {
    if (i == robot.size())
      return 0;
    if (j == factory.length)
      return Long.MAX_VALUE;
    if (mem[i][j][k] > 0)
      return mem[i][j][k];
    final long skipFactory = minimumTotalDistance(robot, factory, i, j + 1, 0, mem);
    final int position = factory[j][0];
    final int limit = factory[j][1];
    final long useFactory = limit > k ? minimumTotalDistance(robot, factory, i + 1, j, k + 1, mem) +
                                            Math.abs(robot.get(i) - position)
                                      : Long.MAX_VALUE / 2;
    return mem[i][j][k] = Math.min(skipFactory, useFactory);
  }
}

class Solution:
  def minimumTotalDistance(
      self,
      robot: list[int],
      factory: list[list[int]],
  ) -> int:
    robot.sort()
    factory.sort()

    @functools.lru_cache(None)
    def dp(i: int, j: int, k: int) -> int:
      """
      Returns the minimum distance to fix robot[i..n) with factory[j..n), where
      factory[j] already fixed k robots.
      """
      if i == len(robot):
        return 0
      if j == len(factory):
        return math.inf
      skipFactory = dp(i, j + 1, 0)
      position, limit = factory[j]
      useFactory = (dp(i + 1, j, k + 1) + abs(robot[i] - position)
                    if limit > k else math.inf)
      return min(skipFactory, useFactory)

    return dp(0, 0, 0)