#880445 (Java) No.732 3PrimeCounting

• 問題: No.732 3PrimeCounting
• ユーザー: ゆうき
• 提出日時: 2023-06-14 02:51:49
• 言語: Java (openjdk 23)
• 結果: AC
• 実行時間: 826 ms / 3,000 ms
• コード長: 9,610 bytes
• コンパイル時間: 4,690 ms
• コンパイル使用メモリ: 99,584 KB
• 実行使用メモリ: 77,316 KB
• 最終ジャッジ日時: 2024-06-22 14:51:39
• 合計ジャッジ時間: 25,296 ms
• ジャッジサーバーID （参考情報）: judge2 / judge5

ファイルパターン	結果
other	AC * 89

ソースコード

import java.io.*;
import java.lang.reflect.Array;
import java.util.*;
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.*;
import java.util.stream.*;
import java.awt.Point;

import static java.lang.Math.*;
import static java.util.Arrays.*;
import static java.util.Collections.*;

class Solver{
  long st = System.currentTimeMillis();

  long elapsed(){ return System.currentTimeMillis() -st; }

  void reset(){ st = System.currentTimeMillis(); }

  static int infI = (int) 1e9;
  static long infL = (long) 1e18;
  static long mod = (int) 1e9 +7;
  //  static long mod = 998244353;
  static String yes = "Yes";
  static String no = "No";

  Random rd = ThreadLocalRandom.current();
  MyReader in = new MyReader(System.in);
  MyWriter out = new MyWriter(System.out);
  MyWriter log = new MyWriter(System.err){
    @Override
    protected void ln(){
      super.ln();
      flush();
    };
  };

  int N = in.it();

  Object solve(){

    int[] PN = primes(N).stream().toArray();
    BitSet primes = primes(3 *N);
    int[] P3N = primes.stream().toArray();

    long[] A = new long[N +1];
    for (var l:PN)
      A[l] = 1;
    long[] X = FastFourierTransform.multiply(A,A);
    X = FastFourierTransform.multiply(X,A);

    long ans = 0;

    for (var p:P3N)
      ans += X[p];
    for (var a:PN)
      for (var b:PN)
        if (primes.get(a *2 +b))
          ans -= 3;

    return ans /6;
  }

  BitSet primes(int b){
    BitSet ret = new BitSet(b);
    ret.set(2,b +1);
    for (int p = 2;p *p <= b;p++)
      if (ret.get(p))
        for (int n = 2 *p;n <= b;n += p)
          ret.clear(n);

    return ret;
  }
}

class FastFourierTransform{
  static void fft(double[] a,double[] b,boolean invert){
    int count = a.length;
    for (int i = 1,j = 0;i < count;i++) {
      int bit = count >>1;
      for (;j >= bit;bit >>= 1)
        j -= bit;
      j += bit;
      if (i < j) {
        double temp = a[i];
        a[i] = a[j];
        a[j] = temp;
        temp = b[i];
        b[i] = b[j];
        b[j] = temp;
      }
    }
    for (int len = 2;len <= count;len <<= 1) {
      int halfLen = len >>1;
      double angle = 2 *Math.PI /len;
      if (invert)
        angle = -angle;
      double wLenA = Math.cos(angle);
      double wLenB = Math.sin(angle);
      for (int i = 0;i < count;i += len) {
        double wA = 1;
        double wB = 0;
        for (int j = 0;j < halfLen;j++) {
          double uA = a[i +j];
          double uB = b[i +j];
          double vA = a[i +j +halfLen] *wA -b[i +j +halfLen] *wB;
          double vB = a[i +j +halfLen] *wB +b[i +j +halfLen] *wA;
          a[i +j] = uA +vA;
          b[i +j] = uB +vB;
          a[i +j +halfLen] = uA -vA;
          b[i +j +halfLen] = uB -vB;
          double nextWA = wA *wLenA -wB *wLenB;
          wB = wA *wLenB +wB *wLenA;
          wA = nextWA;
        }
      }
    }
    if (invert)
      for (int i = 0;i < count;i++) {
        a[i] /= count;
        b[i] /= count;
      }
  }

  static long[] multiply(long[] a,long[] b){
    int resultSize = Integer.highestOneBit(Math.max(a.length,b.length) -1) <<2;
    resultSize = Math.max(resultSize,1);
    double[] aReal = new double[resultSize];
    double[] aImaginary = new double[resultSize];
    double[] bReal = new double[resultSize];
    double[] bImaginary = new double[resultSize];
    for (int i = 0;i < a.length;i++)
      aReal[i] = a[i];
    for (int i = 0;i < b.length;i++)
      bReal[i] = b[i];
    fft(aReal,aImaginary,false);
    if (a == b) {
      System.arraycopy(aReal,0,bReal,0,aReal.length);
      System.arraycopy(aImaginary,0,bImaginary,0,aImaginary.length);
    } else
      fft(bReal,bImaginary,false);
    for (int i = 0;i < resultSize;i++) {
      double real = aReal[i] *bReal[i] -aImaginary[i] *bImaginary[i];
      aImaginary[i] = aImaginary[i] *bReal[i] +bImaginary[i] *aReal[i];
      aReal[i] = real;
    }
    fft(aReal,aImaginary,true);
    long[] result = new long[resultSize];
    for (int i = 0;i < resultSize;i++)
      result[i] = Math.round(aReal[i]);
    return result;
  }

}

class Edge{
  int id;
  int u;
  int v;
  long l;
  Edge rev;

  Edge(int id,int u,int v){
    this.id = id;
    this.u = u;
    this.v = v;
  }

  void rev(Edge rev){ rev.l = l; }

}

class MyReader{
  byte[] buf = new byte[1 <<16];
  int ptr = 0;
  int tail = 0;
  InputStream in;

  MyReader(InputStream in){ this.in = in; }

  byte read(){
    if (ptr == tail)
      try {
        tail = in.read(buf);
        ptr = 0;
      } catch (IOException e) {}
    return buf[ptr++];
  }

  boolean isPrintable(byte c){ return 32 < c && c < 127; }

  boolean isNum(byte c){ return 47 < c && c < 58; }

  byte nextPrintable(){
    byte ret = read();
    while (!isPrintable(ret))
      ret = read();
    return ret;
  }

  int[] arrI(int N,IntUnaryOperator f){
    int[] ret = new int[N];
    setAll(ret,f);
    return ret;
  }

  long[] arrL(int N,IntToLongFunction f){
    long[] ret = new long[N];
    setAll(ret,f);
    return ret;
  }

  double[] arrD(int N,IntToDoubleFunction f){
    double[] ret = new double[N];
    setAll(ret,f);
    return ret;
  }

  <T> T[] arr(T[] arr,IntFunction<T> f){
    setAll(arr,f);
    return arr;
  }

  int it(){ return toIntExact(lg()); }

  int[] it(int N){ return arrI(N,i -> it()); }

  int[][] it(int H,int W){ return arr(new int[H][],i -> it(W)); }

  int idx(){ return it() -1; }

  int[] idx(int N){ return arrI(N,i -> idx()); }

  int[][] idx(int H,int W){ return arr(new int[H][],i -> idx(W)); }

  int[][] qry(int Q){ return arr(new int[Q][],i -> new int[]{idx(), idx(), i}); }

  long lg(){
    byte i = nextPrintable();
    boolean negative = i == 45;
    long n = negative ? 0 : i -'0';
    while (isPrintable(i = read()))
      n = 10 *n +i -'0';
    return negative ? -n : n;
  }

  long[] lg(int N){ return arrL(N,i -> lg()); }

  long[][] lg(int H,int W){ return arr(new long[H][],i -> lg(W)); }

  double dbl(){ return Double.parseDouble(str()); }

  double[] dbl(int N){ return arrD(N,i -> dbl()); }

  double[][] dbl(int H,int W){ return arr(new double[H][],i -> dbl(W)); }

  char[] ch(){ return str().toCharArray(); }

  char[][] ch(int H){ return arr(new char[H][],i -> ch()); }

  String line(){
    StringBuilder sb = new StringBuilder();

    for (byte c;(c = read()) != '\n';)
      sb.append((char) c);
    return sb.toString();
  }

  String str(){
    StringBuilder sb = new StringBuilder();
    sb.append((char) nextPrintable());

    for (byte c;isPrintable(c = read());)
      sb.append((char) c);
    return sb.toString();
  }

  String[] str(int N){ return arr(new String[N],i -> str()); }

  Edge[] e(int N,int M){ return e(N,M,e -> e.l = 1); }

  Edge[] e(int N,int M,Consumer<Edge> f){
    return arr(new Edge[M],i -> {
      Edge e = new Edge(i,idx(),idx());
      f.accept(e);
      return e;
    });
  }

  Edge[][] g(int N,int M,boolean b){ return g(N,b,e(N,M)); }

  Edge[][] g(int N,int M,boolean b,Consumer<Edge> f){ return g(N,b,e(N,M,f)); }

  Edge[][] g(int N,boolean b,Edge[] E){
    int[] c = new int[N];

    for (Edge e:E) {
      c[e.u]++;
      if (!b)
        c[e.v]++;
    }

    Edge[][] g = arr(new Edge[N][],i -> new Edge[c[i]]);

    for (Edge e:E) {

      g[e.u][--c[e.u]] = e;
      if (!b) {
        Edge rev = new Edge(e.id,e.v,e.u);
        e.rev(rev);
        g[e.v][--c[e.v]] = e.rev = rev;
      }
    }
    return g;
  }
}

class MyWriter{
  OutputStream out;
  byte[] buf = new byte[1 <<16];
  byte[] ibuf = new byte[20];
  int tail = 0;

  MyWriter(OutputStream out){ this.out = out; }

  void flush(){
    try {
      out.write(buf,0,tail);
      tail = 0;
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

  protected void ln(){ write((byte) '\n'); }

  private void write(byte b){
    buf[tail++] = b;
    if (tail == buf.length)
      flush();
  }

  private void write(byte[] b,int off,int len){
    for (int i = off;i < off +len;i++)
      write(b[i]);
  }

  private void write(long n){
    if (n < 0) {
      n = -n;
      write((byte) '-');
    }

    int i = ibuf.length;
    do {
      ibuf[--i] = (byte) (n %10 +'0');
      n /= 10;
    } while (n > 0);

    write(ibuf,i,ibuf.length -i);
  }

  private void print(Object obj){
    if (obj instanceof Boolean)
      print((boolean) obj ? Solver.yes : Solver.no);
    else if (obj.getClass().isArray()) {
      int l = Array.getLength(obj);
      for (int i = 0;i < l;i++) {
        print(Array.get(obj,i));
        if (i +1 < l)
          write((byte) ' ');
      }
    } else if (obj instanceof Character)
      write((byte) (char) obj);
    else if (obj instanceof Integer)
      write((int) obj);
    else if (obj instanceof Long)
      write((long) obj);
    else if (obj instanceof char[])
      for (char b:(char[]) obj)
        write((byte) b);
    else
      for (char b:Objects.toString(obj).toCharArray())
        write((byte) b);
  }

  void println(Object obj){
    if (obj == null)
      return;

    if (obj instanceof Collection<?>)
      for (var e:(Collection<?>) obj)
        println(e);
    else if (obj.getClass().isArray()
        && !(Array.get(obj,0) instanceof char[])
        && Array.get(obj,0).getClass().isArray()) {
      int l = Array.getLength(obj);
      for (int i = 0;i < l;i++)
        println(Array.get(obj,i));
    } else {
      print(obj);
      ln();
    }
  }
}

class Main{
  public static void main(String[] args) throws Exception{
    Solver solver = new Solver();
    Optional.ofNullable(solver.solve()).ifPresent(solver.out::println);
    solver.out.flush();
    solver.log.println(solver.elapsed());
  }
}