Skip to content

3279. Maximum Total Area Occupied by Pistons

  • 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
class Solution {
 public:
  long long maxArea(int height, vector<int>& positions, string directions) {
    long area = accumulate(positions.begin(), positions.end(), 0L);
    long ans = area;
    long diffPerSecond = 0;
    map<int, vector<int>> timeToIndices;

    for (int i = 0; i < positions.size(); ++i)
      if (directions[i] == 'U') {
        timeToIndices[height - positions[i]].push_back(i);
        timeToIndices[height - positions[i] + height].push_back(i);
        ++diffPerSecond;
      } else {
        timeToIndices[positions[i]].push_back(i);
        timeToIndices[positions[i] + height].push_back(i);
        --diffPerSecond;
      }

    int prevTime = 0;

    for (const auto& [time, indices] : timeToIndices) {
      area += (time - prevTime) * diffPerSecond;
      ans = max(ans, area);
      prevTime = time;
      for (const int i : indices)
        if (directions[i] == 'U') {
          directions[i] = 'D';
          diffPerSecond -= 2;
        } else {
          directions[i] = 'U';
          diffPerSecond += 2;
        }
    }

    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
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Solution {
  public long maxArea(int height, int[] positions, String directions) {
    long area = Arrays.stream(positions).asLongStream().sum();
    long ans = area;
    long diffPerSecond = 0;
    TreeMap<Integer, List<Integer>> timeToIndices = new TreeMap<>();

    for (int i = 0; i < positions.length; ++i) {
      if (directions.charAt(i) == 'U') {
        timeToIndices.computeIfAbsent(height - positions[i], k -> new ArrayList<>()).add(i);
        timeToIndices.computeIfAbsent(height - positions[i] + height, k -> new ArrayList<>())
            .add(i);
        ++diffPerSecond;
      } else {
        timeToIndices.computeIfAbsent(positions[i], k -> new ArrayList<>()).add(i);
        timeToIndices.computeIfAbsent(positions[i] + height, k -> new ArrayList<>()).add(i);
        --diffPerSecond;
      }
    }

    char[] directionsChars = directions.toCharArray();
    int prevTime = 0;

    for (Map.Entry<Integer, List<Integer>> entry : timeToIndices.entrySet()) {
      final int time = entry.getKey();
      List<Integer> indices = entry.getValue();
      area += (time - prevTime) * diffPerSecond;
      ans = Math.max(ans, area);
      prevTime = time;
      for (final int i : indices)
        if (directionsChars[i] == 'U') {
          directionsChars[i] = 'D';
          diffPerSecond -= 2;
        } else {
          directionsChars[i] = 'U';
          diffPerSecond += 2;
        }
    }

    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
25
26
27
28
29
30
31
32
33
34
35
36
from sortedcontainers import SortedDict


class Solution:
  def maxArea(self, height: int, positions: list[int], directions: str) -> int:
    area = sum(positions)
    ans = area
    diffPerSecond = 0
    timeToIndices: SortedDict[int, list[int]] = SortedDict()

    for i, (position, direction) in enumerate(zip(positions, directions)):
      if direction == 'U':
        timeToIndices.setdefault(height - position, []).append(i)
        timeToIndices.setdefault(height - position + height, []).append(i)
        diffPerSecond += 1
      else:
        timeToIndices.setdefault(position, []).append(i)
        timeToIndices.setdefault(position + height, []).append(i)
        diffPerSecond -= 1

    prevTime = 0
    directionsList = list(directions)

    for time, indices in timeToIndices.items():
      area += (time - prevTime) * diffPerSecond
      ans = max(ans, area)
      prevTime = time
      for i in indices:
        if directionsList[i] == 'U':
          directionsList[i] = 'D'
          diffPerSecond -= 2
        else:
          directionsList[i] = 'U'
          diffPerSecond += 2

    return ans