#848440 (Java21) No.206 数の積集合を求めるクエリ

提出ソース

結果

問題	No.206 数の積集合を求めるクエリ
ユーザー	ゆうき
提出日時	2023-03-11 17:16:43
言語	Java21 (openjdk 21)
結果	AC
実行時間	414 ms / 7,000 ms
コード長	11,695 bytes
コンパイル時間	4,838 ms
コンパイル使用メモリ	94,704 KB
実行使用メモリ	64,972 KB
最終ジャッジ日時	2024-09-18 06:36:57
合計ジャッジ時間	11,131 ms
ジャッジサーバーID （参考情報）	judge4 / judge1

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	62 ms 50,288 KB
testcase_01	AC	65 ms 50,608 KB
testcase_02	AC	64 ms 50,748 KB
testcase_03	AC	64 ms 50,584 KB
testcase_04	AC	64 ms 50,600 KB
testcase_05	AC	66 ms 50,576 KB
testcase_06	AC	101 ms 52,292 KB
testcase_07	AC	94 ms 51,940 KB
testcase_08	AC	95 ms 52,000 KB
testcase_09	AC	99 ms 52,168 KB
testcase_10	AC	62 ms 50,616 KB
testcase_11	AC	65 ms 50,392 KB
testcase_12	AC	99 ms 52,384 KB
testcase_13	AC	101 ms 52,280 KB
testcase_14	AC	102 ms 52,424 KB
testcase_15	AC	100 ms 52,020 KB
testcase_16	AC	98 ms 52,164 KB
testcase_17	AC	278 ms 64,616 KB
testcase_18	AC	293 ms 63,084 KB
testcase_19	AC	335 ms 64,956 KB
testcase_20	AC	291 ms 63,020 KB
testcase_21	AC	278 ms 63,164 KB
testcase_22	AC	265 ms 63,240 KB
testcase_23	AC	339 ms 64,652 KB
testcase_24	AC	414 ms 64,972 KB
testcase_25	AC	389 ms 64,908 KB
testcase_26	AC	337 ms 63,100 KB
testcase_27	AC	337 ms 63,148 KB
testcase_28	AC	367 ms 63,224 KB
testcase_29	AC	365 ms 63,240 KB
testcase_30	AC	343 ms 63,020 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

import java.io.*;
import java.util.*;
import java.util.concurrent.ThreadLocalRandom;
import java.util.function.*;

class Solver{
  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();
  long st = System.currentTimeMillis();

  static InputStream is;
  MyReader in = new MyReader(is);
  MyWriter out = new MyWriter(System.out);
  MyWriter log = new MyWriter(System.err){
    @Override
    void ln(){
      super.ln();
      flush();
    };
  };

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

  int L = in.it();
  int M = in.it();
  int N = in.it();
  int[] A = in.it(L);
  int[] B = in.it(M);
  int Q = in.it();

  Object solve(){
    long[] a = new long[N +1];
    long[] b = new long[N +1];
    for (var aa:A)
      a[aa]++;
    for (var bb:B)
      b[N -bb]++;

    long[] ans = FastFourierTransform.multiply(a,b);

    for (int q = 1;q <= Q;q++)
      out.println(ans[N +q -1]);
    return null;
  }

}

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 SegmentTree<V> {
  Seg<V, V> seg;

  SegmentTree(Seg<V, V> seg){ this.seg = seg; }

  void build(Function<Integer, V> sup){
    seg.build(sup);
    for (int i = seg.n;--i > 0;)
      seg.merge(i);
  }

  void upd(int i,V f){
    seg.upd(i,i +1,f);
    seg.up(i,i +1);
  }

  V get(int i){ return get(i,i +1); }

  V get(int l,int r){ return seg.get(l,r); }
}

class DualSegmentTree<V, F> {
  Seg<V, F> seg;

  DualSegmentTree(Seg<V, F> seg){ this.seg = seg; }

  void upd(int i,F f){ upd(i,i +1,f); }

  void upd(int l,int r,F f){
    seg.down(l,r);
    seg.upd(l,r,f);
  }

  V get(int i){
    seg.down(i,i +1);
    return seg.get(i,i +1);
  }

}

class LazySegmentTree<V, F> {
  Seg<V, F> seg;

  LazySegmentTree(Seg<V, F> seg){ this.seg = seg; }

  void build(Function<Integer, V> sup){
    seg.build(sup);
    for (int i = seg.n;--i > 0;)
      seg.merge(i);
  }

  void upd(int i,F f){ upd(i,i +1,f); }

  void upd(int l,int r,F f){
    seg.down(l,r);
    seg.upd(l,r,f);
    seg.up(l,r);
  }

  V get(int i){ return get(i,i +1); }

  V get(int l,int r){
    seg.down(l,r);
    return seg.get(l,r);
  }
}

abstract class Seg<V, F> {
  int n;
  V e;
  V[] val;
  F[] lazy;
  int[][] lr;

  @SuppressWarnings("unchecked")
  Seg(int n,V e){
    this.n = n;
    this.e = e;
    val = (V[]) new Object[n <<1];
    lazy = (F[]) new Object[n];

    lr = new int[n <<1][];
    for (int i = n <<1;--i > 0;)
      lr[i] = new int[]{i < n ? lr[i <<1][0] : i, i < n ? lr[i <<1 |1][1] : i +1};
  }

  void build(Function<Integer, V> sup){
    for (int i = 0;i < n;i++)
      val[i +n] = sup.apply(i);
  }

  abstract V agg(V v0,V v1);

  abstract V map(V v,F f);

  abstract F comp(F f0,F f1);

  abstract F powF(F f,int[] lr);

  void merge(int i){ val[i] = agg(eval(i <<1),eval(i <<1 |1)); }

  void uprec(int l,int r){
    merge(l >>1);
    if (l != r)
      merge(r >>1);
    if (0 < l && 0 < r)
      uprec(l >>1,r >>1);
    else if (0 < l)
      uprec(l >>1,r);
    else if (0 < r)
      uprec(l,r >>1);

  }

  void up(int l,int r){
    l += n;
    r += n;
    uprec(l /(l &-l),r /(r &-r));
  }

  void comp(int i,F f){
    if (i < n)
      lazy[i] = lazy[i] != null ? comp(lazy[i],f) : f;
    else
      val[i] = map(val[i],f);
  }

  V eval(int i){

    if (i < n && lazy[i] != null) {
      val[i] = map(val[i],powF(lazy[i],lr[i]));
      for (int c = 0;c < 2;c++)
        comp(i <<1 |c,lazy[i]);
      lazy[i] = null;
    }

    return val[i] != null ? val[i] : e;
  }

  void downrec(int l,int r){
    if (0 < l && 0 < r)
      downrec(l >>1,r >>1);
    else if (0 < l)
      downrec(l >>1,r);
    else if (0 < r)
      downrec(l,r >>1);

    eval(l);
    if (l != r)
      eval(r);
  }

  void down(int l,int r){
    l += n;
    r += n;
    downrec(l /(l &-l),r /(r &-r));
  }

  void upd(int l,int r,F f){
    l += n;
    r += n;
    do {
      if ((l &1) == 1)
        comp(l++,f);
      if ((r &1) == 1)
        comp(--r,f);
    } while ((l >>= 1) < (r >>= 1));
  }

  V get(int l,int r){
    if (l >= r)
      return null;
    l += n;
    r += n;
    V vl = e;
    V vr = e;
    do {
      if ((l &1) == 1)
        vl = agg(vl,eval(l++));
      if ((r &1) == 1)
        vr = agg(eval(--r),vr);
    } while ((l >>= 1) < (r >>= 1));

    return agg(vl,vr);
  }

}

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

  static boolean local() throws Exception{
    Solver.is = new FileInputStream("src/input.txt");
    return true;
  }
}

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 it(){ return (int) lg(); }

  int[] it(int N){
    int[] a = new int[N];
    Arrays.setAll(a,i -> it());
    return a;
  }

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

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

  int[] idx(int N){
    int[] a = new int[N];
    Arrays.setAll(a,i -> idx());
    return a;
  }

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

  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){
    long[] a = new long[N];
    Arrays.setAll(a,i -> lg());
    return a;
  }

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

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

  double[] dbl(int N){
    double[] a = new double[N];
    Arrays.setAll(a,i -> dbl());
    return a;
  }

  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;isPrintable(c = read()) || c == ' ';)
      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()); }

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

  List<Set<Integer>> g(int N,int M,boolean b){
    List<Set<Integer>> g = new ArrayList<>();
    for (int i = 0;i < N;i++)
      g.add(new HashSet<>());
    while (M-- > 0) {
      int u = idx();
      int v = idx();
      g.get(u).add(v);
      if (!b)
        g.get(v).add(u);
    }
    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();
    }
  }

  void sp(){ write((byte) ' '); }

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

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

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

  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);
  }

  void println(boolean b){ println(b ? Solver.yes : Solver.no); }

  void println(long n){
    write(n);
    ln();
  }

  void println(double d){ println(String.valueOf(d)); }

  void println(String s){ println(s.toCharArray()); }

  void println(char[] s){
    for (char b:s)
      write((byte) b);
    ln();
  }

  void println(int[] a){
    for (int i = 0;i < a.length;i++) {
      if (0 < i)
        sp();
      write(a[i]);
    }
    ln();
  }

  void println(long[] a){
    for (int i = 0;i < a.length;i++) {
      if (0 < i)
        sp();
      write(a[i]);
    }
    ln();
  }

  void println(double[] a){
    for (int i = 0;i < a.length;i++) {
      if (0 < i)
        sp();
      for (char b:String.valueOf(a[i]).toCharArray())
        write((byte) b);
    }
    ln();
  }

  void println(Object obj){
    if (obj instanceof Boolean)
      println((boolean) obj);
    else if (obj instanceof char[])
      println((char[]) obj);
    else if (obj instanceof int[])
      println((int[]) obj);
    else if (obj instanceof long[])
      println((long[]) obj);
    else if (obj instanceof double[])
      println((double[]) obj);
    else
      println(Objects.toString(obj));
  }
}

yukicoder

結果

テストケース

ソースコード