Four Arrays FOURARR Solution Codechef

Four Arrays FOURARR Solution Codechef

You have been given an integer K and four arrays A, B, C, and D with sizes SA, SB, SC and SD respectively.

Follow the given steps:

Pick four indices x, y, z, and w from arrays A, B, C, and D respectively. Note that (1≤x≤SA), (1≤y≤SB), (1≤z≤SC) and (1≤w≤SD).
There are in total SA×SB×SC×SD ways to do the above step. For each way, write down the value of (Ax+By)×(Cz+Dw) on a board.
Sort all the numbers written on the board in non-decreasing order.
Find the Kth element written on the board.

Input Format

  • First line contains five integers SA, SB, SC, SD and K.
  • The second line contains SA space-separated integers A1,A2,…,ASA denoting the array A.
  • The third line contains SB space-separated integers B1,B2,…,BSB denoting the array B.
  • The fourth line contains SC space-separated integers C1,C2,…,CSC denoting the array C.
  • The fifth line contains SD space-separated integers D1,D2,…,DSD denoting the array D.

Output Format
For each test case, output the required number in a single line .

Constraints

  • 1≤SA,SB,SC,SD≤3⋅104
  • 1≤K≤SA×SB×SC×SD
  • 0≤Ai,Bi,Ci,Di≤105

Sample Input 1
2 1 3 2 10
2 3
2
1 1 2
3 1

Sample Output 1
20

Explanation
After sorting the numbers on the board are in the following order {8,8,10,10,12,15,16,16,20,20,20,25}. So, the 10th number is 20.

SOLUTION

Program: Four Arrays FOURARR Solution in C++

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;

typedef tree<
int,
null_type,
less<int>,
rb_tree_tag,
tree_order_statistics_node_update>
ordered_set;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const int mod = 1e9+7;

typedef complex<double> CD;
//struct CD {
//    double x, y;
//    CD(double x=0, double y=0) :x(x), y(y) {}
//    CD operator+(const CD& o) { return {x+o.x, y+o.y};}
//    CD operator-(const CD& o) { return {x-o.x, y-o.y};}
//    CD operator*(const CD& o) { return {x*o.x-y*o.y, x*o.y+o.x*y};}
//    void operator /= (double d) { x/=d; y/=d;}
//    double real() {return x;}
//    double imag() {return y;}
//};
//CD conj(const CD &c) {return CD(c.x, -c.y);}

const double PI = acos(-1.0L);

namespace FFT {
    int N;
    vector<int> perm;
    vector<CD> wp[2];

    void precalculate(int n) {
        assert((n & (n-1)) == 0);
        N = n;
        perm = vector<int> (N, 0);
        for (int k=1; k<N; k<<=1) {
            for (int i=0; i<k; i++) {
                perm[i] <<= 1;
                perm[i+k] = 1 + perm[i];
            }
        }

        wp[0] = wp[1] = vector<CD>(N);
        for (int i=0; i<N; i++) {
            wp[0][i] = CD( cos(2*PI*i/N),  sin(2*PI*i/N) );
            wp[1][i] = CD( cos(2*PI*i/N), -sin(2*PI*i/N) );
        }
    }

    void fft(vector<CD> &v, bool invert = false) {
        if (v.size() != perm.size())    precalculate(v.size());

        for (int i=0; i<N; i++)
            if (i < perm[i])
                swap(v[i], v[perm[i]]);

        for (int len = 2; len <= N; len *= 2) {
            for (int i=0, d = N/len; i<N; i+=len) {
                for (int j=0, idx=0; j<len/2; j++, idx += d) {
                    CD x = v[i+j];
                    CD y = wp[invert][idx]*v[i+j+len/2];
                    v[i+j] = x+y;
                    v[i+j+len/2] = x-y;
                }
            }
        }

        if (invert) {
            for (int i=0; i<N; i++) v[i]/=N;
        }
    }

     void pairfft(vector<CD> &a, vector<CD> &b, bool invert = false) {
        int N = a.size();
        vector<CD> p(N);
        for (int i=0; i<N; i++) p[i] = a[i] + b[i] * CD(0, 1);
        fft(p, invert);
        p.push_back(p[0]);

        for (int i=0; i<N; i++) {
            if (invert) {
                a[i] = CD(p[i].real(), 0);
                b[i] = CD(p[i].imag(), 0);
            }
            else {
                a[i] = (p[i]+conj(p[N-i]))*CD(0.5, 0);
                b[i] = (p[i]-conj(p[N-i]))*CD(0, -0.5);
            }
        }
    }

    vector<ll> multiply(const vector<ll> &a, const vector<ll> &b) {
        int n = 1;
        while (n < a.size()+ b.size())  n<<=1;

        vector<CD> fa(a.begin(), a.end()), fb(b.begin(), b.end());
        fa.resize(n); fb.resize(n);

//        fft(fa); fft(fb);
        pairfft(fa, fb);
        for (int i=0; i<n; i++) fa[i] = fa[i] * fb[i];
        fft(fa, true);

        vector<ll> ans(n);
        for (int i=0; i<n; i++)     ans[i] = round(fa[i].real());
        return ans;
    }

    const int M = 1e9+7, B = sqrt(M)+1;
    vector<ll> anyMod(const vector<ll> &a, const vector<ll> &b) {
        int n = 1;
        while (n < a.size()+ b.size())  n<<=1;
        vector<CD> al(n), ar(n), bl(n), br(n);

        for (int i=0; i<a.size(); i++)  al[i] = a[i]%M/B, ar[i] = a[i]%M%B;
        for (int i=0; i<b.size(); i++)  bl[i] = b[i]%M/B, br[i] = b[i]%M%B;

        pairfft(al, ar); pairfft(bl, br);
//        fft(al); fft(ar); fft(bl); fft(br);

        for (int i=0; i<n; i++) {
            CD ll = (al[i] * bl[i]), lr = (al[i] * br[i]);
            CD rl = (ar[i] * bl[i]), rr = (ar[i] * br[i]);
            al[i] = ll; ar[i] = lr;
            bl[i] = rl; br[i] = rr;
        }

        pairfft(al, ar, true); pairfft(bl, br, true);
//        fft(al, true); fft(ar, true); fft(bl, true); fft(br, true);

        vector<ll> ans(n);
        for (int i=0; i<n; i++) {
            ll right = round(br[i].real()), left = round(al[i].real());;
            ll mid = round(round(bl[i].real()) + round(ar[i].real()));
            ans[i] = ((left%M)*B*B + (mid%M)*B + right)%M;
        }
        return ans;
    }
}

void test_case(){
    int na,nb,nc,nd;
    ll k; cin >> na >> nb >> nc >> nd >> k;
    vector<int> a(na),b(nb),c(nc),d(nd);
    vector<ll> ha(100001,0),hb(100001,0),hc(100001,0),hd(100001,0);
    for(int i=0;i<na;i++){
        cin >> a[i];
        ha[a[i]]++;
    }
    for(int i=0;i<nb;i++){
        cin >> b[i];
        hb[b[i]]++;
    }
    for(int i=0;i<nc;i++){
        cin >> c[i];
        hc[c[i]]++;
    }
    for(int i=0;i<nd;i++){
        cin >> d[i];
        hd[d[i]]++;
    }

    vector<ll> fi = FFT::multiply(ha,hb), se = FFT::multiply(hc,hd);
    fi.resize(200001), se.resize(200001);

    ll l = -1, r = 40000000005;
    while(l+1<r){
        ll mid = (l+r)>>1;

        // count number of elements less than or equal to mid
        ll ans = 0;
        int j = 200000;
        int cnt = nc*nd;
        for(int i=0;i<=200000;i++){
            while(1ll*i*j>mid){
                cnt -= se[j];
                --j;
            }
            ans += 1ll*cnt*fi[i];
        }

        if(ans >= k){
            r = mid;
        }
        else{
            l = mid;
        }
    }

    cout << r << endl;
    return;
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t = 1;
    //cin >> t;
    while(t--){
        test_case();
    }
    return 0;
}

Program: Four Arrays FOURARR Solution in Java

import java.io.OutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.PrintWriter;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.AbstractCollection;
import java.io.InputStream;

/**
 * Built using CHelper plug-in
 * Actual solution is at the top
 *
 * @author sarthakmanna
 */
public class Main {
    public static void main(String[] args) {
        InputStream inputStream = System.in;
        OutputStream outputStream = System.out;
        FastReader in = new FastReader(inputStream);
        PrintWriter out = new PrintWriter(outputStream);
        FourArrays solver = new FourArrays();
        solver.solve(1, in, out);
        out.close();
    }

    static class FourArrays {
        final int MAXN = 200_005;
        int[][] A;
        long K;
        int[] sumAB;
        int[] sumCD;
        FenwickTree fenCD;

        public void fillInputParams(FastReader in) {
            int[] s = in.getIntArray(4);
            K = in.nextLong();
            A = new int[4][MAXN / 2];
            for (int idx = 0; idx < 4; ++idx) {
                for (int i = 0; i < s[idx]; ++i) ++A[idx][in.nextInt()];
            }
        }

        Object solveOptimised(FastReader in, StringBuilder sb) {
            sumAB = multiply(A[0], A[1]);
            sumCD = multiply(A[2], A[3]);
            fenCD = new FenwickTree(MAXN);
            for (int i = 0; i < MAXN; ++i) fenCD.update(i, sumCD[i]);

            long l = 0, r = Long.MAX_VALUE >> 1;
            while (true) {
                long mid = l + r >> 1;
                if (l == mid) {
                    if (countFloors(l) >= K) return l;
                    else return r;
                } else if (countFloors(mid) < K) {
                    l = mid + 1;
                } else {
                    r = mid;
                }
            }
        }

        long countFloors(long val) {
            long count = sumAB[0] * fenCD.prefixQuery(MAXN - 1);

            for (int i = 1; i < MAXN; ++i) {
                // sumAB * sumCD <= val
                int remSum = (int) Math.min(MAXN - 1, val / i);
                count += sumAB[i] * fenCD.prefixQuery(remSum);
            }

            return count;
        }

        int[] multiply(int[] a, int[] b) {
            double[] aDouble = new double[a.length], bDouble = new double[b.length];
            for (int i = 0; i < a.length; ++i) aDouble[i] = a[i];
            for (int i = 0; i < b.length; ++i) bDouble[i] = b[i];

            FFT fft = new FFT();
            double[] prod = fft.multiply(aDouble, bDouble);

            int[] c = new int[MAXN];
            for (int i = 0; i < Math.min(MAXN, prod.length); ++i) {
                c[i] = (int) Math.round(prod[i]);
            }
            return c;
        }

        Object solveBrute(FastReader in, StringBuilder sb) {
            return null;
        }

        public void solve(int testNumber, FastReader in, PrintWriter out) {
            // out.print("Case #" + testNumber + ": ");

            fillInputParams(in);
            Object outOptimised = solveOptimised(in, new StringBuilder());
            Object outBrute = solveBrute(in, new StringBuilder());
            if (outBrute == null) {
                out.println(outOptimised);
            } else if (outBrute.toString().equals(outOptimised.toString())) {
                System.err.println(testNumber + ". OK Checked");
            } else {
                // print input params
                System.err.println("Actual = " + outOptimised);
                System.err.println("Expected = " + outBrute);
                System.err.println();
                throw new ArithmeticException();
            }
        }

    }

    static class FastReader {
        static final int BUFSIZE = 1 << 20;
        static byte[] buf;
        static int index;
        static int total;
        static InputStream in;

        public FastReader(InputStream is) {
            try {
                in = is;
                buf = new byte[BUFSIZE];
            } catch (Exception e) {
            }
        }

        private int scan() {
            try {
                if (index >= total) {
                    index = 0;
                    total = in.read(buf);
                    if (total <= 0)
                        return -1;
                }
                return buf[index++];
            } catch (Exception | Error e) {
                System.err.println(e.getMessage());
                return 13 / 0;
            }
        }

        public int nextInt() {
            int c, val = 0;
            for (c = scan(); c <= 32; c = scan()) ;
            boolean neg = c == '-';
            if (c == '-' || c == '+')
                c = scan();
            for (; c >= '0' && c <= '9'; c = scan())
                val = (val << 3) + (val << 1) + (c & 15);
            return neg ? -val : val;
        }

        public long nextLong() {
            int c;
            long val = 0;
            for (c = scan(); c <= 32; c = scan()) ;
            boolean neg = c == '-';
            if (c == '-' || c == '+')
                c = scan();
            for (; c >= '0' && c <= '9'; c = scan())
                val = (val << 3) + (val << 1) + (c & 15);
            return neg ? -val : val;
        }

        public int[] getIntArray(int size) {
            int[] ar = new int[size];
            for (int i = 0; i < size; ++i) ar[i] = nextInt();
            return ar;
        }

    }

    static class FenwickTree {
        int N;
        long[] tree;

        public FenwickTree(int size) {
            N = size + 1;
            tree = new long[N];
        }

        public void update(int idx, long val) {
            ++idx;
            while (idx < N) {
                tree[idx] += val;
                idx += idx & -idx;
            }
        }

        public long prefixQuery(int idx) {
            ++idx;
            long ret = 0;
            while (idx > 0) {
                ret += tree[idx];
                idx -= idx & -idx;
            }
            return ret;
        }

        public long query(int l, int r) {
            if (l > r) return 0;
            long ret = prefixQuery(r);
            if (l > 0) ret -= prefixQuery(l - 1);
            return ret;
        }

        public String toString() {
            ArrayList<Long> toStr = new ArrayList<>();
            for (int i = 0; i < N - 1; ++i) toStr.add(query(i, i));
            return toStr.toString();
        }

    }

    static class FFT {
        public 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;
                }
            }
        }

        public double[] multiply(double[] a, double[] 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);
            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);
            return aReal;
        }

    }
}

Related:

Leave a Comment

three × one =