[FIXED] How can I make this Java program run faster?

Issue

I am solving this problem from HackerEarth but I am getting the "Time limit exceeded" error.

Function F(N) is defined as:

F(N) = sum_of_digits(N²)

Given a number N, output if it is possible for it to ultimately become {1 or 4} or not.

Input:
First line contains T which is the number of test cases.
T lines follow each with an integer N.

Output:
For each N output "YES" or "NO" if the number can achieve desired goal or not.

This is the program:

import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.stream.Stream;

public class MoodyNumbers {

    public static boolean is1or4(int x) {
        boolean result = false;
        Integer digitSum = Stream.of((int) Math.pow(x, 2)).reduce(Integer::sum).get();
        StringBuilder sumStringBuilder = new StringBuilder(String.valueOf(digitSum));
        for (int i = 0; i < sumStringBuilder.length(); i++) {
            result = sumStringBuilder.charAt(i) == '1' || sumStringBuilder.charAt(i) == '4';
            if (result) {
                break;
            }
        }
        return result;
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedOutputStream bos = new BufferedOutputStream(System.out);
        int numberOfTest = Integer.parseInt(br.readLine());
        for (int i = 0; i < numberOfTest; i++) {
            boolean result = is1or4(Integer.parseInt(br.readLine()));
            if (result) {
                bos.write("YES\n".getBytes());
            }
            else {
                bos.write("NO\n".getBytes());
            }
            bos.flush();
        }
    }
}

I ran out of ideas to make this faster, any recommendations?

Solution

What takes time is writing to standard output, so if you write the result to standard output for each test case, that will increase the overall time significantly.

You can use StringBuilder and append to it the result of each test case. Then after you have handled all the test cases, print out the entire contents of the StringBuilder.

Also, I think you misunderstood the question. Your code checks whether the sum of the digits of the square of the test case number (i.e. the result of function F(N)) contains a 1 (one) or a 4. As I understand it, you need to check whether the result of F(N) is 1 or 4.

Here is my solution which was accepted on hackerearth, i.e. it produced the correct results in an acceptable time.

import java.util.Scanner;

public class MoodyNumbers {

    private static boolean is1or4(int x) {
        boolean result = x == 1 || x == 4;
        if (!result) {
            int newX = sumDigits((long) x * x);
            if (newX == 1 || newX == 4) {
                return true;
            }
            if (newX == 9 || newX == 16) {
                return false;
            }
            result = is1or4(newX);
        }
        return result;
    }

    private static int sumDigits(long x) {
        int sum = 0;
        while (x >= 10) {
            sum += x % 10;
            x /= 10;
        }
        sum += x;
        return sum;
    }

    public static void main(String[] args) {
        Scanner stdin = new Scanner(System.in);
        int numberOfTestCases = stdin.nextInt();
        StringBuilder sb = new StringBuilder(numberOfTestCases * 4);
        String newLine = System.lineSeparator();
        while (numberOfTestCases-- > 0) {
            int n = stdin.nextInt();
            sb.append((is1or4(n) ? "YES" : "NO"));
            sb.append(newLine);
        }
        System.out.print(sb);
    }
}

If the result of F(N) is either 9 or 16, that creates an endless loop. F(9) returns 9. F(16) returns 13 and F(13) returns 16 which also creates an endless loop. Hence for any number for which F(N) eventually returns either 9 or 16, the result is NO.

Answered By - Abra
Answer Checked By - David Marino (JavaFixing Volunteer)