Kitty's Calculations on a Tree

  • + 0 comments

    I'm not sure if we can get a solution in js/ts. I have adapted a working c++ solution. It gets tests 1-6 correct, 7-9 wrong, everything else times out.

    Code:

    'use strict';
    
    process.stdin.resume();
    process.stdin.setEncoding('utf-8');
    let inputString: string = '';
    let inputLines: string[]=[];
    process.stdin.on('data', function(inputStdin: string): void {
        inputString += inputStdin;
    });
    
    function readLine():string{
        return inputLines.shift();
    }
    
    const K = 1000000007;
    const L = 200001;
    
    let parents : number[] = Array(L);
    let children : number[] = Array(L);
    function resetChildren(){
        children = Array(L).fill(0);
    }
    let valuesSet : boolean[] = Array(L-1);
    function resetValuesSet(){
        valuesSet= Array(L-1).fill(false);
    }
    
    process.stdin.on('end', function(): void {
        inputLines = inputString.split('\n');
        inputString = '';
        main();
    });
    
    
    function main(){
        const [n,q] = initMain();
        // edges
        prepareEdges(n);
        queryloop:for(let i = 0; i < q; i++){
            const k = parseInt(readLine());
            if(k < 2) {
                // result = 0
                process.stdout.write(`0\n`);
                // skip next line
                readLine();
                // continue loop
                continue queryloop;
            }
            const query = readLine().split(" ").map(v => parseInt(v));
        
            // reset valueSet
            resetValuesSet();
    
            // reset children
            resetChildren();
            
            let valuesSum = 0;
            for (const q of query) {
                valuesSum += q;
                valuesSet[q] = true;
            }
            let sum = 0;
            for (let j = n; j > 0; j--) {
                let a = children[j];
                if (valuesSet[j]) {
                    a += j;
                }
                if (a) {
                    let x = a * (valuesSum - a);// ((a % K) * ((valuesSum - a) % K)) % K; //
                    // sum =(sum + x)%K;
                    sum += x;
                    // if(sum>K) sum = sum % K;
                }
                children[parents[j]] += a;
            }
            process.stdout.write(`${sum%K}\n`);
        }
    }
    
    function initMain(){
        return readLine().split(" ").map(v => parseInt(v));
    }
    
    function prepareEdges(n:number){
        for (let i = 1; i < n; i++){
            const [a,b] = readLine().split(" ").map(v=>parseInt(v));
            if(a<b)
                parents[b] = a;
            else
                parents[a] = b;
        }
        parents[1]=0;
    }