백준 17499: 수열과 시프트 쿼리

https://www.acmicpc.net/problem/17499

문제 이해

길이가 $N$ 인 정수 수열 $[a_1, a_2, ..., a_N]$ 이 주어집니다.
이 수열에 적용되는 연산은 세 가지는 다음과 같습니다.
- 1 i x: $a_i$ 에 정수 x만큼 더하기
- 2 s: 수열을 오른쪽으로 s칸 시프트
- 3 s: 수열을 왼쪽으로 s칸 시프트

1차 시도

접근 방법

반복문이 N번 실행되어 너무 느리다는 문제의 문구를 잘못 해석해서 배열이 아니라 LinkedList로 구현하면 문제를 해결할 수 있지 않을까? 라는 생각을 했습니다.

코드

public class Main {

  public static void main(String[] args) throws IOException {
    BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    StringBuilder sb = new StringBuilder();
    StringTokenizer st = new StringTokenizer(br.readLine());

    // 문제 입력
    int N = Integer.parseInt(st.nextToken());
    int Q = Integer.parseInt(st.nextToken());

    LinkedList<Integer> q = new LinkedList<>();
    Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).forEach(q::offer);

    // 쿼리 시작
    for (int i = 0; i < Q; i++) {
      st = new StringTokenizer(br.readLine());
      String op = st.nextToken();
      int a = Integer.parseInt(st.nextToken());

      switch (op) {
        case "1":
          // a_i에 x만큼 더함
          int b = Integer.parseInt(st.nextToken());
          q.set(a - 1, q.get(a - 1) + b);
          break;
        case "2":
          // 수열을 오른쪽으로 a칸 시프트
          for (int j = a; j > 0; j--) {
            q.offerFirst(q.pollLast());
          }
          break;
        case "3":
          // 수열을 왼쪽으로 a칸 시프트
          for (int j = a; j > 0; j--) {
            q.offerLast(q.pollFirst());
          }
          break;
      }
    }

    q.forEach(n -> sb.append(n).append(" "));
    System.out.println(sb.toString().trim());
  }

}

2차 시도

실패 원인 분석 & 해결 방안

총 수열의 길이 N은 최대 200,000이고, 시프트 할 칸 s는 N-1 이하입니다. 때문에 반복문이 N번 실행되어 너무 느리다는 것은 어떤 테스트 케이스에서 20만 번의 시프트가 연산의 개수만큼 반복될 수 있다는 뜻이 됩니다.
때문에 이 문제를 해결하기 위해서는 시프트 연산에 시프트를 하지 말아야 합니다. ~~(네...?)~~
현재 수열의 시작 위치를 저장할 변수를 생성하고 이를 시프트 반대 방향으로 움직입니다.
- 예를 들어, 문제 예시에 있듯 수열 $[a_1, a_2, ..., a_N]$ 을 왼쪽으로 한 칸 시프트할 때 왼쪽에서 pop하고 오른쪽에 push하여 $[a_2, ..., a_N, a_1]$ 를 만드는 대신
- 최초 시작 $a_1$ 을 가리키던 변수 int p = 0을 $a_2$ 를 가리키도록 p = 1로 수정해 줍니다.
- 이 방법을 쓴다면 10만 번 시프트 하라는 쿼리가 입력되더라도 10만 번의 push, pop 연산이 아니라 p += 100000 연산 한 번이면 해결됩니다.
- 다만 이 경우 p의 값이 0 ~ N 범위를 벗어나지 않도록 잘 관리해야 합니다.

코드

public class Main {

  public static void main(String[] args) throws IOException {
    BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    StringBuilder sb = new StringBuilder();
    StringTokenizer st = new StringTokenizer(br.readLine());

    // 문제 입력
    int N = Integer.parseInt(st.nextToken());
    int Q = Integer.parseInt(st.nextToken());

    int[] arr = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();

    // 수열 시작을 가리킬 인덱스
    int p = 0;

    // 쿼리 시작
    for (int i = 0; i < Q; i++) {
      st = new StringTokenizer(br.readLine());
      String op = st.nextToken();
      int a = Integer.parseInt(st.nextToken());

      switch (op) {
        case "1":
          // a_i에 x만큼 더함
          int b = Integer.parseInt(st.nextToken());
          int aa = (N + p + (a - 1)) % N;
          arr[aa] += b;
          break;
        case "2":
          // 수열을 오른쪽으로 a칸 시프트
          // -> 시작 인덱스를 왼쪽으로 a칸 옮김
          // 인덱스 범위 주의
          p = (N + p - a) % N;
          break;
        case "3":
          // 수열을 왼쪽으로 a칸 시프트
          // -> 시작 인덱스를 오른쪽으로 a칸 옮김
          // 인덱스 범위 주의
          p = (N + p + a) % N;
          break;
      }
    }

    // 수열 시작을 가리키는 인덱스를 기준으로
    // 결과 문자열로 합치기
    for (int i = p1; i < N; i++) {
      sb.append(arr[i]).append(' ');
    }
    for (int i = 0; i < p1; i++) {
      sb.append(arr[i]).append(' ');
    }

    System.out.println(sb.toString().trim());
  }

}

결과

92852KB / 836ms

시간이 조금 오래 걸린 것 같긴 하지만, 시간 초과가 발생하던 첫 번째 시도와는 달리 제한 시간 내에 통과할 수 있었습니다. 문제에서 친절하게 설명해주는 것을 잘 읽을 필요가 있고, LinkedList 같은 자료구조가 항상 적합하지는 않다는 것을 기억할 필요가 있습니다. (LinkedList는 느리다, 틀렸다가 아니고 이 문제에 적합하지 않았다는 것입니다.)

Prerequisites: 이 글에서 언급되었으나 깊게 설명하지 않은 내용입니다.
Java LinkedList
Java Queue

P.S. 이 문제의 이름에는 '수열'과 '쿼리'가 들어갑니다. 그래서 solved.ac 계정이 있는 분들은 이 문제를 풀면 이 멋진 수열에 쿼리를! 배경을 받을 수 있습니다.