Codeforces Round #378 (Div. 2) E. Sleep in Class
問題
http://codeforces.com/contest/733/problem/E
N 階建ての建物の中にいます。各階には 'U' または 'D' の文字が書かれています。 i 階にいて U と書かれているときは i+1 階へ、D と書かれているときは i-1 階へ移動します。移動すると、直前にいた階の文字は切り替わります。i 回にいるとき、何回移動すると N+1 階もしくは 0 階へ移動できるか出力してください。
解法
実は、状況は問題文ほどややこしくありません。右側の一番近くのDで跳ね返り、左側の一番近くのUで跳ね返り、右側の二番目に近いDで跳ね返り…を外に出るまで繰り返します。累積和を上手に使うと O(N) で N 個分を一気に求めることができます。
コード
import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; import java.io.PrintWriter; import java.util.NoSuchElementException; /* _ooOoo_ o8888888o 88" . "88 (| -_- |) O\ = /O ____/`---'\____ .' \\| |// `. / \\||| : |||// \ / _||||| -:- |||||- \ | | \\\ - /// | | | \_| ''\---/'' | | \ .-\__ `-` ___/-. / ___`. .' /--.--\ `. . __ ."" '< `.___\_<|>_/___.' >'"". | | : `- \`.;`\ _ /`;.`/ - ` : | | \ \ `-. \_ __\ /__ _/ .-` / / ======`-.____`-.___\_____/___.-`____.-'====== `=---=' ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ pass System Test! */ public class Main { private static class Task { void solve(FastScanner in, PrintWriter out) throws Exception { int N = in.nextInt(); char[] S = in.next().toCharArray(); int[] leftUp = new int[N]; for (int i = 1; i < N; i++) { leftUp[i] = leftUp[i - 1]; if (S[i - 1] == 'U') leftUp[i]++; } int[] rightDown = new int[N]; for (int i = N - 2; i >= 0; i--) { rightDown[i] = rightDown[i + 1]; if (S[i + 1] == 'D') rightDown[i]++; } long[] leftWing = new long[N]; long[] rightWing = new long[N]; { long rightSum = 0; int cur = 0; int right = 0; for (int i = 0; i < N; i++) { rightSum -= right * 2; if (S[i] == 'D' && right > 0) right--; int left = leftUp[i]; if (S[i] == 'U') left++; if (cur <= i) cur = i + 1; while (cur < N && right < left) { if (S[cur] == 'D') { right++; rightSum += (cur - i) * 2; } cur++; } rightWing[i] = rightSum; } } { long leftSum = 0; int cur = N - 1; int left = 0; for (int i = N - 1; i >= 0; i--) { leftSum -= left * 2; if (S[i] == 'U' && left > 0) left--; int right = rightDown[i]; if (S[i] == 'D') right++; if (i <= cur) cur = i - 1; while (cur >= 0 && left < right) { if (S[cur] == 'U') { left++; leftSum += (i - cur) * 2; } cur--; } leftWing[i] = leftSum; } } for (int i = 0; i < N; i++) { long ans = leftWing[i] + rightWing[i]; if (S[i] == 'U') leftUp[i]++; else rightDown[i]++; if (leftUp[i] < rightDown[i] || (leftUp[i] == rightDown[i] && S[i] == 'U')) ans += (i + 1); if (leftUp[i] > rightDown[i] || (leftUp[i] == rightDown[i] && S[i] == 'D')) ans += (N - i); out.print(ans + " "); } out.println(); } } /** * ここから下はテンプレートです。 */ public static void main(String[] args) throws Exception { OutputStream outputStream = System.out; FastScanner in = new FastScanner(); PrintWriter out = new PrintWriter(outputStream); Task solver = new Task(); solver.solve(in, out); out.close(); } private static class FastScanner { private final InputStream in = System.in; private final byte[] buffer = new byte[1024]; private int ptr = 0; private int bufferLength = 0; private boolean hasNextByte() { if (ptr < bufferLength) { return true; } else { ptr = 0; try { bufferLength = in.read(buffer); } catch (IOException e) { e.printStackTrace(); } if (bufferLength <= 0) { return false; } } return true; } private int readByte() { if (hasNextByte()) return buffer[ptr++]; else return -1; } private static boolean isPrintableChar(int c) { return 33 <= c && c <= 126; } private void skipUnprintable() { while (hasNextByte() && !isPrintableChar(buffer[ptr])) ptr++; } boolean hasNext() { skipUnprintable(); return hasNextByte(); } public String next() { if (!hasNext()) throw new NoSuchElementException(); StringBuilder sb = new StringBuilder(); int b = readByte(); while (isPrintableChar(b)) { sb.appendCodePoint(b); b = readByte(); } return sb.toString(); } long nextLong() { if (!hasNext()) throw new NoSuchElementException(); long n = 0; boolean minus = false; int b = readByte(); if (b == '-') { minus = true; b = readByte(); } if (b < '0' || '9' < b) { throw new NumberFormatException(); } while (true) { if ('0' <= b && b <= '9') { n *= 10; n += b - '0'; } else if (b == -1 || !isPrintableChar(b)) { return minus ? -n : n; } else { throw new NumberFormatException(); } b = readByte(); } } double nextDouble() { return Double.parseDouble(next()); } double[] nextDoubleArray(int n) { double[] array = new double[n]; for (int i = 0; i < n; i++) { array[i] = nextDouble(); } return array; } double[][] nextDoubleMap(int n, int m) { double[][] map = new double[n][]; for (int i = 0; i < n; i++) { map[i] = nextDoubleArray(m); } return map; } public int nextInt() { return (int) nextLong(); } public int[] nextIntArray(int n) { int[] array = new int[n]; for (int i = 0; i < n; i++) array[i] = nextInt(); return array; } public long[] nextLongArray(int n) { long[] array = new long[n]; for (int i = 0; i < n; i++) array[i] = nextLong(); return array; } public String[] nextStringArray(int n) { String[] array = new String[n]; for (int i = 0; i < n; i++) array[i] = next(); return array; } public char[][] nextCharMap(int n) { char[][] array = new char[n][]; for (int i = 0; i < n; i++) array[i] = next().toCharArray(); return array; } public int[][] nextIntMap(int n, int m) { int[][] map = new int[n][]; for (int i = 0; i < n; i++) { map[i] = nextIntArray(m); } return map; } } }