[PAT-A] 1128 N Queens Puzzle

Java

测试点3超时

import java.io.*;

public class Main {
    static final int N = 1010;

    static boolean checkConflict(int[] queens, int col) {
        for (int i = 1; i < col; i++) {
            if (queens[i] == queens[col] || Math.abs(queens[i] - queens[col]) == col - i)
                return true;
        }
        return false;
    }

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

        int K = Integer.parseInt(br.readLine());

        while (K-- > 0) {
            String[] inputs = br.readLine().split(" ");
            int N = Integer.parseInt(inputs[0]);

            int[] queens = new int[N + 1];
            boolean valid = true;
            for (int i = 1; i <= N; i++) {
                queens[i] = Integer.parseInt(inputs[i]);
                if (checkConflict(queens, i)) {
                    valid = false;
                    break;
                }
            }
            System.out.println(valid ? "YES" : "NO");
        }

        br.close();
    }
}

C++

#include <iostream>
#include <cstring>

using namespace std;

const int N = 1010;

int n;
bool col[N], diag[N * 2], rdiag[N * 2];        // diag=y-x, rdiag=y+x

int main() {
    cin >> n;
    while (n--) {
        memset(col, false, sizeof col);
        memset(diag, false, sizeof diag);
        memset(rdiag, false, sizeof rdiag);

        int m;
        cin >> m;

        bool flag = true;
        for (int x = 1; x <= m; ++x) {
            int y;
            cin >> y;

            if (col[y] || diag[y - x + m] || rdiag[y + x])
                flag = false;
            col[y] = diag[y - x + m] = rdiag[y + x] = true;
        }

        if (flag) cout << "YES" << endl;
        else cout << "NO" << endl;
    }

    return 0;
}