Submission #3582794

---

Source Code Expand

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.*;

public class Main {

    static int N, M;
    static int X0, A, P;

    public static void main(String[] args) {
        FastScanner sc = new FastScanner(System.in);
        N = sc.nextInt();
        M = sc.nextInt();
        X0 = sc.nextInt();
        A = sc.nextInt();
        P = sc.nextInt();

        System.out.println(solve());
    }

    static long solve() {
        Student[] S = new Student[N*M];
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < M; j++) {
                int idx = i*M+j;
                int x = (int)((X0 + ((long)idx*A)) % P);
                S[idx] = new Student(i, j, x);
            }
        }
        Arrays.sort(S, Comparator.comparingInt(s -> s.x));

        boolean[] row = new boolean[M];
        PriorityQueue<Student> dq = new PriorityQueue<>( Comparator.comparingInt(s -> s.j) );
        long ans = 0;
        for (int i = 0; i < N; i++) {
            Arrays.fill(row, false);
            int rowAns = 0;

            for (int j = 0; j < M; j++) {
                Student s = S[i*M+j];
                if( s.i == i ) {
                    row[s.j] = true;
                } else {
                    rowAns += Math.abs(s.i - i);
                    dq.add(s);
                }
            }

            for (int j = 0; j < M; j++) {
                if( !row[j] ) {
                    Student s = dq.poll();
                    rowAns += Math.abs(s.j - j);
                }
            }

            ans += rowAns;
        }
        return ans;
    }

    static class Student {
        int i, j; // 元の位置
        int x; // 成績

        public Student(int i, int j, int x) {
            this.i = i;
            this.j = j;
            this.x = x;
        }

        int idx() {
            return i * M + j;
        }
    }

    @SuppressWarnings("unused")
    static class FastScanner {
        private BufferedReader reader;
        private StringTokenizer tokenizer;

        FastScanner(InputStream in) {
            reader = new BufferedReader(new InputStreamReader(in));
            tokenizer = null;
        }

        String next() {
            if (tokenizer == null || !tokenizer.hasMoreTokens()) {
                try {
                    tokenizer = new StringTokenizer(reader.readLine());
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
            return tokenizer.nextToken();
        }

        String nextLine() {
            if (tokenizer == null || !tokenizer.hasMoreTokens()) {
                try {
                    return reader.readLine();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
            return tokenizer.nextToken("\n");
        }

        long nextLong() {
            return Long.parseLong(next());
        }

        int nextInt() {
            return Integer.parseInt(next());
        }

        int[] nextIntArray(int n) {
            int[] a = new int[n];
            for (int i = 0; i < n; i++)
                a[i] = nextInt();
            return a;
        }

        long[] nextLongArray(int n) {
            long[] a = new long[n];
            for (int i = 0; i < n; i++)
                a[i] = nextLong();
            return a;
        }
    }
}

Submission Info

Judge Result