如何优化Java程序解决HackerEarth Moody Numbers超时问题？

核心问题分析 #

你的代码存在两个关键问题：

• 逻辑错误：仅计算一次F(N)（即N²的数位和）就判断，没有按照题目要求反复计算F直到结果为1或4。
• 效率低下：使用Stream、字符串转换等低效操作计算数位和，频繁IO flush，且无结果缓存。

具体优化措施 #

1. 修正循环计算逻辑
   必须反复计算F，直到结果为1/4，或检测到循环（说明无法到达目标值）。

2. 高效计算数位和
   用数学取模+除法替代字符串和Stream操作，避免不必要的对象创建和开销：

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

3. 避免浮点数运算
   用(long)x * x替代Math.pow(x, 2)，既避免浮点数精度丢失，又提升运算速度。

4. 预先缓存结果
   由于F(N)的结果范围极小（最大数位和仅为162），预先计算所有可能值的结果并缓存，每个测试用例直接查表：

   private static boolean[] cache = new boolean[200];
   
   static {
       for (int i = 1; i < 200; i++) {
           cache[i] = canReach1Or4(i);
       }
   }
   
   private static boolean canReach1Or4(int x) {
       boolean[] visited = new boolean[200];
       while (true) {
           if (x == 1 || x == 4) return true;
           if (visited[x]) return false;
           visited[x] = true;
           long square = (long) x * x;
           x = sumDigits(square);
       }
   }
   

5. 优化IO操作
   用StringBuilder批量收集所有输出，最后一次性写入，避免频繁flush的IO开销：

   public static void main(String[] args) throws IOException {
       BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
       BufferedOutputStream bos = new BufferedOutputStream(System.out);
       int t = Integer.parseInt(br.readLine());
       StringBuilder sb = new StringBuilder();
       for (int i = 0; i < t; i++) {
           long n = Long.parseLong(br.readLine());
           int f = sumDigits(n * n);
           sb.append(cache[f] ? "YES\n" : "NO\n");
       }
       bos.write(sb.toString().getBytes());
       bos.flush();
       bos.close();
       br.close();
   }
   

完整优化后代码 #

import java.io.BufferedOutputStream;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;

public class MoodyNumbers {
    private static boolean[] cache = new boolean[200];

    static {
        for (int i = 1; i < 200; i++) {
            cache[i] = canReach1Or4(i);
        }
    }

    private static boolean canReach1Or4(int x) {
        boolean[] visited = new boolean[200];
        while (true) {
            if (x == 1 || x == 4) {
                return true;
            }
            if (visited[x]) {
                return false;
            }
            visited[x] = true;
            long square = (long) x * x;
            x = sumDigits(square);
        }
    }

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

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedOutputStream bos = new BufferedOutputStream(System.out);
        int t = Integer.parseInt(br.readLine());
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < t; i++) {
            long n = Long.parseLong(br.readLine());
            int f = sumDigits(n * n);
            sb.append(cache[f] ? "YES\n" : "NO\n");
        }
        bos.write(sb.toString().getBytes());
        bos.flush();
        bos.close();
        br.close();
    }
}

内容的提问来源于stack exchange，提问作者Santiagolp