getMonthlyBalances函数月度股票余额计算错误排查与修复求助

问题背景 #

getMonthlyBalances函数用于根据股票交易历史计算月度持仓余额，输入为stocks（含股票ID、月度价格）和trades（含交易记录）数组。当前函数存在计算错误，例如测试数据中2月的余额应为1000，但函数返回0。

原函数代码 #

function getMonthlyBalances(stocks, trades) {
  if (!stocks || !trades) {
    return [];
  }

  if (trades.length === 0) {
    return {};
  }

  // Sort trades by timestamp
  trades.sort((a, b) => a.timestamp.seconds - b.timestamp.seconds);

  // Get the first trade month and create a date object
  const firstTradeDate = new Date(trades[0].timestamp.seconds * 1000);
  firstTradeDate.setDate(1); // Set the date to the first day of the month
  firstTradeDate.setHours(0, 0, 0, 0); // Set the time to midnight

  const monthlyBalances = {};
  let currentDate = firstTradeDate;

  // Initialize the user's stock holdings
  const userStocks = {};

  while (currentDate <= new Date()) {
    const currentMonth = currentDate.getMonth();
    const currentYear = currentDate.getFullYear();
    const currentMonthStr = `${currentYear}-${currentMonth + 1}`;

    trades.forEach(trade => {
      const tradeDate = new Date(trade.timestamp.seconds * 1000);
      if (tradeDate.getFullYear() === currentYear && tradeDate.getMonth() === currentMonth) {
        if (trade.type === "buy") {
          userStocks[trade.id] = (userStocks[trade.id] || 0) + trade.quantity;
        } else if (trade.type === "sell") {
          userStocks[trade.id] = (userStocks[trade.id] || 0) - trade.quantity;
        }
      }
    });

    let balance = 0;
    for (const stockId in userStocks) {
      const stock = stocks.find(stock => stock.id === parseInt(stockId));
      const stockPrice = stock.prices.find(price => {
        const priceDate = new Date(price.date.seconds * 1000);
        return priceDate.getFullYear() === currentYear && priceDate.getMonth() === currentMonth;
      });

      if (stockPrice) {
        balance += userStocks[stockId] * stockPrice.price;
      }
    }

    monthlyBalances[currentMonthStr] = balance;

    currentDate.setMonth(currentDate.getMonth() + 1);
  }

  return monthlyBalances;
}

输入示例数据 #

// trades
[
  {
    "id": 1, // stock_1
    "quantity": 10,
    "price": 120,
    "type": "buy",
    "timestamp": {
      "seconds": 1672589479,
      "nanoseconds": 858000000
    }
  },
  {
    "id": 1, // stock_1
    "quantity": 10,
    "price": 110,
    "type": "buy",
    "timestamp": {
      "seconds": 1675354636,
      "nanoseconds": 858000000
    }
  }
]

// stocks
[
  {
    "id": 1,
    "prices": [
      {
        "price": 120,
        "date": {
          "seconds": 1678860000,
          "nanoseconds": 267000000
        }
      },
      {
        "price": 100,
        "date": {
          "seconds": 1672552800,
          "nanoseconds": 32000000
        }
      }
    ]
  },
  {
    "id": 2,
    "prices": [
      {
        "price": 120,
        "date": {
          "seconds": 1677650400,
          "nanoseconds": 12000000
        }
      },
      {
        "price": 90,
        "date": {
          "seconds": 1672552800,
          "nanoseconds": 145000000
        }
      }
    ]
  }
]

错误原因分析 #

1. 价格匹配逻辑缺失：函数仅查找与当前月份完全匹配的股票价格，若当月无对应价格则直接忽略该持仓的价值。例如测试数据中，stock1仅在1月和3月有价格，2月无数据，导致2月余额计算时这部分持仓未被计入，返回0，但实际应沿用最近的1月价格（100）计算10股的价值（1000）。
2. 交易处理效率低下且易出错：每个月份循环时都遍历所有交易，判断是否属于当前月后更新持仓。虽当前逻辑不会重复处理交易，但冗余遍历会降低效率，且后续修改逻辑时容易引入重复计算的bug。
3. 返回值类型不一致：当输入为空时，函数可能返回数组（[]）或对象（{}），不符合一致的输出规范。

修复方案 #

1. 改进价格匹配逻辑：为每个股票的价格按时间排序，查找小于等于当前月份的最近价格，确保无当月价格时能沿用历史价格。
2. 优化交易处理流程：维护一个交易指针，随着月份推进，逐步处理到当前月份为止的所有未处理交易，避免重复遍历全部交易。
3. 统一返回值类型：所有空输入场景均返回空对象，与正常输出格式一致。

修复后代码 #

function getMonthlyBalances(stocks, trades) {
  // 统一空输入返回空对象
  if (!stocks || !trades || trades.length === 0) {
    return {};
  }

  // 按时间排序交易
  trades.sort((a, b) => a.timestamp.seconds - b.timestamp.seconds);
  // 预处理股票：按时间排序价格，方便后续查找最近价格
  const processedStocks = stocks.map(stock => ({
    ...stock,
    prices: [...stock.prices].sort((a, b) => a.date.seconds - b.date.seconds)
  }));

  const firstTradeDate = new Date(trades[0].timestamp.seconds * 1000);
  firstTradeDate.setDate(1);
  firstTradeDate.setHours(0, 0, 0, 0);

  const monthlyBalances = {};
  let currentDate = firstTradeDate;
  const userStocks = {};
  let tradeIndex = 0; // 交易处理指针，避免重复遍历

  while (currentDate <= new Date()) {
    const currentMonth = currentDate.getMonth();
    const currentYear = currentDate.getFullYear();
    const currentMonthStr = `${currentYear}-${currentMonth + 1}`;
    const currentMonthEnd = new Date(currentYear, currentMonth + 1, 0);

    // 处理当前月份及之前未处理的交易
    while (tradeIndex < trades.length) {
      const trade = trades[tradeIndex];
      const tradeDate = new Date(trade.timestamp.seconds * 1000);
      if (tradeDate > currentMonthEnd) {
        break; // 交易时间晚于当前月份，停止处理
      }
      // 更新持仓
      if (trade.type === "buy") {
        userStocks[trade.id] = (userStocks[trade.id] || 0) + trade.quantity;
      } else if (trade.type === "sell") {
        userStocks[trade.id] = (userStocks[trade.id] || 0) - trade.quantity;
      }
      tradeIndex++;
    }

    let balance = 0;
    for (const stockId in userStocks) {
      const stockIdNum = parseInt(stockId);
      const stock = processedStocks.find(s => s.id === stockIdNum);
      if (!stock) continue;

      // 查找最近的、不晚于当前月份的价格
      let latestPrice = null;
      for (const price of stock.prices) {
        const priceDate = new Date(price.date.seconds * 1000);
        if (priceDate.getFullYear() < currentYear || 
            (priceDate.getFullYear() === currentYear && priceDate.getMonth() <= currentMonth)) {
          latestPrice = price;
        } else {
          break; // 价格已按时间排序，后续价格更晚，无需继续
        }
      }

      if (latestPrice) {
        balance += userStocks[stockId] * latestPrice.price;
      }
    }

    monthlyBalances[currentMonthStr] = balance;
    // 推进到下一个月
    currentDate.setMonth(currentDate.getMonth() + 1);
  }

  return monthlyBalances;
}

验证结果 #

使用示例输入运行修复后的函数，将得到符合期望的输出：

{
  "2023-1": 1000,
  "2023-2": 1000,
  "2023-3": 2400
}

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